from collections import OrderedDict
import math

import numpy as np
import pandas as pd
import scipy

from mobmess import utils
from mobmess.dummy_module import DummyModule

try:
    import matplotlib.pyplot as plt
except ImportError:
    plt = DummyModule("matplotlib")

try:
    import seaborn as sns
except ImportError:
    sns = DummyModule("seaborn")

def histplot(*args,
             xlabel=None, ylabel=None, title=None, show=None,
             **kwargs):
    """
    Wrapper around seaborn.histplot. Adds extra convenience parameters
    """

    ax = sns.histplot(*args, **kwargs)
    ax.set(xlabel=xlabel, ylabel=ylabel, title=title)

    if show:
        plt.show()

    return ax

def linkage(data, metric='precompute', method='ward', fast=False, **kwargs):
    """Does the same thing with scipy.cluster.hierarchy.linkage, but it
    handles my desired default parameters, and easily handles the
    usage of a precomputed distance metric.

    fast :

        If True, use fastcluster package's implementation of linkage(). Otherwise, use scipy's (default).

        Note: not recommended to use fastcluster unless necessary. My
        empirical tests show that the fastcluster package is about 2x
        faster with method='ward' and precomputed distances for
        ~16K-by-16K.  However, in a direct comparison of Z created by
        scipy, it is not exactly the same -- although discrepancies
        may just be degenerate permutations of the
        leaves. Additionally, I'm not sure if fastcluster is computing
        optimal leaf ordering (this is done by default in scipy, which
        may also exlain why it's slower)
    """

    if metric=='precompute':
        data = scipy.spatial.distance.squareform(data, checks=False)
    # else:
    #     raise Exception()

    if fast:
        import fastcluster
        linkage = fastcluster.linkage
    else:
        linkage = scipy.cluster.hierarchy.linkage

    # print('metric', metric)
    # print(data)
    Z = linkage(data, metric=metric, method=method, **kwargs)

    return Z

def plot_colors(colors_df, ncols=None):
    """Visualizes the mapping of values to their colors as created by utils.make_category_colors()"""

    if ncols is None:
        ncols = 3

    # Remove unused categorical. These will attempted to be colored (which is wasteful of color space) and plotted (which results in an error)
    colors_df = utils.unused_cat(colors_df)

    # Map values in the dataframe to colors
    # - dataframe with colors
    columns = colors_df.columns
#    print(columns)
    colors_df = pd.concat([colors_df,
                           make_category_colors(None, colors_df, fmt='rgb').rename(columns=lambda x : f'{x}_color')],
                           axis=1)
#    display(colors_df)

    fig, axes = subplots(None, ncols, 3, 1, n=colors_df.shape[1], flatten=True)
    for ax, c in zip(axes, columns):
        data = colors_df[[c, f'{c}_color']].drop_duplicates().copy()
        if len(data)<=1:
            continue

        data['height'] = 1
        # display(data)
        # display(data.dtypes)
        # print(c)
        sns.barplot(x=c, y='height', data=data, ax=ax, palette=data.set_index(c)[f'{c}_color'].to_dict())
        ax.tick_params(labelleft=False, left=False, labelrotation=-45, direction='in')
        ax.set(xlabel=None, ylabel=None)
        for x in ['left', 'bottom', 'top', 'right']:
            ax.spines[x].set_visible(False)
#        ax.set_title(c)
        ax.set_ylabel(c)
    for ax in axes[len(columns) : ]:
        ax.remove()

def make_category_colors(names, assignments, palette_name='husl', fmt='hex', na_color=None):
    """Given a categorization of a set of objects (`names`), assign a
    color to each object by their category. Multiple types of
    categorizations are given as a dataframe

    assignments : dataframe of objects-by-categories

    names : the subset of object to be colored

    na_color : the color used for objects with no category (value of nan in `assignments`)

    """

    import seaborn as sns

    if isinstance(assignments, pd.Series):
        assignments = assignments.to_frame()

    if names is None:
        assignments = assignments.astype('category')
    else:
        assignments = assignments.loc[names].astype('category')

    colors = pd.DataFrame(index=assignments.index)
    n_classes_dict = OrderedDict()
    for c in assignments.columns:
        assignments[c] = assignments[c].cat.remove_unused_categories()
        n_classes = assignments[c].cat.categories.size
        n_classes_dict[c] = n_classes
        palette = sns.color_palette(palette_name, n_classes)

        if fmt=='hex':
            palette = palette.as_hex()
            if na_color is None:
                na_color = '#ffffff'
        elif fmt=='rgb':
            if na_color is None:
                na_color = (1.,1.,1.)

        colors[c] = [palette[code] if (code != -1) else na_color \
                        for code in assignments[c].cat.codes]

    return colors

def subplots(nrows, ncols, width, height, total_height=None, total_width=None,
             flatten=False, width_ratios=None, height_ratios=None, tight_layout=True, n=None, downsize=False, figsize=None, gridspec_kw=None, **kwargs):
    """Convenient wrapper around plt.subplots(). Calculates figsize
    based on (width, height) of each axes"""

    if gridspec_kw is None:
        gridspec_kw = {}

    if width_ratios is not None:
        assert 'width_ratios' not in gridspec_kw
        gridspec_kw['width_ratios'] = width_ratios
    if height_ratios is not None:
        assert 'height_ratios' not in gridspec_kw
        gridspec_kw['height_ratios'] = height_ratios

    if n is not None:
        if ncols is None:
            # Infer number of columns
            ncols = math.ceil(n / nrows)
        else:
            # Infer number of rows
            nrows = math.ceil(n / ncols)

            if downsize and (nrows==1):
                # Decrease number of columns becaues `n` is small
                ncols = min(n, ncols * nrows)

    if figsize is None:
        if total_width is None:
            total_width = ncols * width
        if total_height is None:
            total_height = nrows * height
        figsize = (total_width, total_height)

    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize,
                             tight_layout=tight_layout, squeeze=False, gridspec_kw=gridspec_kw,
                             **kwargs)
    if flatten:
        axes = axes.flatten()

    return fig, axes

def clustermap(data,
               Z=None,
               method=None,
               figsize=None,
               metric=None,
               colors_ratio=None,
               dendrogram_ratio=None,
               is_affinity=True,
               row_colors=None,
               col_colors=None,
               number_labels=False,
               show_cluster=None,
               show_row_cluster=None,
               show_col_cluster=None,
               verbose_colors=False,
               optimal_ordering=None,
               center=None, cmap=None,
               square=None,
               row_Z=None, col_Z=None,
               verbose=True,
               show=False,
               co_cluster=None, n_clusters=None, random_state=None,
               row_cluster=None, col_cluster=None,
               xticklabels=True,
               yticklabels=True,
               cbar_pos=None,
               convert_row_colors=None,
               **kwargs):
    """
    **** DEPRECATED : See clustermap() *************

    Plots a heatmap. Acts like sns.clustermap, but uses my desired
    default paramters, and does basic formatting I typically want

    Notes on runtime:

    ## Seems like biggest time is to plot pixels
    #
    ## Using distances as-is
    # -- 4m 36s, all plasmids, 16827-by-16827
         suggested figsize=(30,30)
    # -- 2m 20s, 12K-by-12K, ward, figsize=20-by-20, row/col_cluster default or False (this setting doesn't matter).
         But you want to let row/col_cluster=True to reorder the leaves
         suggested figsize=(20,20)
    # -- 34s, 6000-by-6000, ward
    # -- 16s, 6000-by-6000, complete
    ## From scratch, default metric
    # -- 24s for 3000-by-3000
    # -- 3m 12s for 6000-by-6000
    """

    import seaborn as sns

    from scipy.cluster import hierarchy

    if optimal_ordering is None:
        optimal_ordering = False

    if square is None:
        square = False

    if metric is None:
        if square:
            metric = 'precompute'
        else:
            metric = 'euclidean'

    if figsize is None:
        if data.shape[0] > 15000:
            figsize = (32,30)
        elif data.shape[0] > 6000:
            figsize = (22,20)
        elif data.shape[0] > 3000:
            figsize = (17, 15)
        elif data.shape[0] > 100:
            figsize = (11,10)
        elif data.shape[0] > 40:
            figsize = (8.75, 8)
        else:
#            figsize = (5.5, 5)
            figsize = (7.5, 7)

    def rearrange_by_Z(data, data_arr, names, Z, row=True):
        """Rearranges the data rows or cols according to a linkage Z"""

        leaves = hierarchy.leaves_list(Z)
        if row:
            # This linkage represents rows
            data, data_arr, names = data.iloc[leaves,:], data_arr[leaves,:], names[leaves]
        else:
            # This linkage represents cols
            data, data_arr, names = data.iloc[:, leaves], data_arr[:,leaves], names[leaves]

        return data, data_arr, names

    if isinstance(data, pd.DataFrame):
        # Test if dataframe contains sparse datatypes. If so, then convert to dense
        if any('Sparse' in x.name for x in data.dtypes.values):
            data = data.sparse.to_dense()

        rownames, colnames = data.index.values, data.columns.values

        if square and not np.all(rownames==colnames):
            assert Z is None

            # Sync rownames and colnames to make square
            data = data.copy()
            union = data.index.union(data.columns)
            data = data.reindex(union).reindex(columns=union).fillna(0)
            rownames, colnames = data.index.values, data.columns.values

        data_arr = data.values
    else:
        rownames, colnames = np.arange(data.shape[0]), np.arange(data.shape[1])
        data_arr = data
        data = pd.DataFrame(data)

    if verbose:
        print('Plotting shape:', data.shape)

    if square:
        assert data.shape[0]==data.shape[1]
        if Z is not None:
            assert (row_Z is not None) and (col_Z is not None)
            row_Z, col_Z = Z, Z

    if is_affinity:
        # The specified matrix is affinities/similarities (rather than
        # distances). So invert it, before giving it to linkage()
        data_linkage = 1 - data_arr
    else:
        data_linkage = data_arr

    if show_cluster:
        assert (show_row_cluster is None) and (show_col_cluster is None)
        show_row_cluster, show_col_cluster = True, True

    to_rearrange_rows, to_rearrange_cols = False, False
    if co_cluster:
        ### Do bi-clustering using sklearn's SpectralCoclustering() method
        if verbose:
            print('Calculating co-clustering')

        import sklearn.cluster
        if random_state is None:
            random_state = 0
        if n_clusters is None:
            n_clusters = 20
        model = sklearn.cluster.SpectralCoclustering(n_clusters=n_clusters, random_state=random_state)
        # print('wacka', np.isinf(data_arr).sum())
        # print('wacka', data_arr.dtype)

        # Normalize values
        data_arr = sklearn.preprocessing.StandardScaler().fit_transform(data_arr)

        data_arr = data_arr.astype(np.float64) + 1e-10 # Add a small number to avoid floating point errors (don't know why this happens)
#        return data_arr
        model.fit(data_arr)
        data = data.iloc[np.argsort(model.row_labels_), :].iloc[:, np.argsort(model.column_labels_)]
        data_arr = data_arr[np.argsort(model.row_labels_), :][:, np.argsort(model.column_labels_)]

        assert (row_Z is None) and (col_Z is None) and (not show_row_cluster) and (not show_col_cluster)
        show_row_cluster, show_col_cluster = False, False
    else:
        ### Cluster rows and cols separately

        # If either row_Z or col_Z is set, then make sure that
        # row/col_cluster is set to the default None or is not False
        if (row_Z is None) and (row_cluster is None):
            # Default is to cluster rows
            row_cluster = True
        elif row_Z is not None:
            if row_cluster is None:
                row_cluster = True
            else:
                assert row_cluster is not False, "If row_Z is specified, then row_cluster cannot be False"
        if (col_Z is None) and (col_cluster is None):
            # Default is to cluster cols
            col_cluster = True
        elif col_Z is not None:
            if col_cluster is None:
                col_cluster = True
            else:
                assert col_cluster is not False, "If col_Z is specified, then col_cluster cannot be False"

        # print(row_cluster, col_cluster)
        # print(show_row_cluster, show_col_cluster)

        if method is None:
            method = 'ward'

        if row_cluster:
            # Make sure a row linkage is calculated
            if row_Z is None:
                if verbose:
                    utils.tprint('Calculating row linkage with metric={}, method={}, optimal_ordering={}'.format(metric, method, optimal_ordering))
                row_Z = linkage(data_linkage, metric=metric, method=method, optimal_ordering=optimal_ordering)

            if show_row_cluster is None:
                # If row_cluster is True, then default is to show them
                show_row_cluster = True

            if show_row_cluster:
                to_rearrange_rows = True
            else:
                ## If not showing the row dendrogram, then need to internally reorder the rows
                data, data_arr, rownames = rearrange_by_Z(data, data_arr, rownames, row_Z, row=True)
        else:
            assert not show_row_cluster, "Cannot show a clustering when one isn't being calculated"
            show_row_cluster = False

        if col_cluster:
            # Make sure a col linkage is calculated
            if col_Z is None:
                if verbose:
                    utils.tprint('Calculating col linkage with metric={}, method={}, optimal_ordering={}'.format(metric, method, optimal_ordering))
                col_Z = linkage(data_linkage.T, metric=metric, method=method, optimal_ordering=optimal_ordering)

            if show_col_cluster is None:
                # If col_cluster is True, then default is to show them
                show_col_cluster = True

            if show_col_cluster:
                to_rearrange_cols = True
            else:
                ## If not showing the col dendrogram, then need to internally reorder the cols
                data, data_arr, colnames = rearrange_by_Z(data, data_arr, colnames, col_Z, row=False)
                to_rearrange_cols = False
        else:
            assert not show_col_cluster, "Cannot show a clustering when one isn't being calculated"
            show_col_cluster = False

    # Calculate colors for heatmap
    if row_colors is not None:
        row_colors = row_colors.loc[row_colors.index.intersection(data.index)]
        row_colors = utils.unused_cat(row_colors)

        if square:
            assert col_colors is None, "Cannot specify `row_colors` and `col_colors` simultaneously when square==True"
            col_colors = row_colors

        if (convert_row_colors is None) or convert_row_colors:
            plot_colors(row_colors)
            row_colors = make_category_colors(None, row_colors, fmt='rgb')

    if col_colors is not None:
        empty_col_colors = col_colors.reindex(colnames).isna().all()
        if verbose_colors:
            print('These col_colors are all NA across rows: {}'.format(empty_col_colors.index[empty_col_colors]))
            display(col_colors.loc[colnames,~empty_col_colors])

        col_colors = col_colors.loc[col_colors.index.intersection(data.columns)]
        col_colors = utils.unused_cat(col_colors)

#        display(col_colors)
        plot_colors(col_colors)
        # print(colnames)
        # display(col_colors)
        # display(empty_col_colors)
        col_colors = make_category_colors(colnames, col_colors.loc[:,~empty_col_colors], fmt='rgb')

    if colors_ratio is None:
        ncol_colors = 1 if (col_colors is None) else col_colors.shape[1]
        nrow_colors = 1 if (row_colors is None) else row_colors.shape[1]
        colors_ratio = (1/figsize[0] / nrow_colors, 1/figsize[1] / ncol_colors)
        #print('colors_ratio:', colors_ratio)

    if dendrogram_ratio is None:
        dendrogram_ratio = (min(0.1, 2 / figsize[0]), min(0.1, 1 / figsize[1]))

    if cbar_pos=='None':
        cbar_pos = None # A workaround to disable the cbar by setting it to 'None' (the string), rather than None
    elif cbar_pos is None:
        # Set the colorbar to a thin vertical bar
        height = min(0.8, 10/figsize[1])
#        cbar_pos=(0, 0.8 - height, min(2/figsize[0], 0.025), height)
        cbar_pos=(0.02, 0.02, min(2/figsize[0], 0.025), height)
        # print('figsize:', figsize)
        # print('cbar_pos:', cbar_pos)

    if verbose:
        utils.tprint('Figure size:', figsize)
        utils.tprint('colors_ratio:', colors_ratio)
        utils.tprint('dendrogram_ratio:', dendrogram_ratio)
        utils.tprint('cbar_pos:', cbar_pos)
        utils.tprint('row/col_cluster:', row_cluster, col_cluster)
        utils.tprint('show_row/col_cluster:', show_row_cluster, show_col_cluster)
        utils.tprint('row/col_Z is set:', row_Z is not None, col_Z is not None)
        utils.tprint('Plotting')

    fig = sns.clustermap(data,
                         row_cluster=show_row_cluster,
                         col_cluster=show_col_cluster,
                         row_linkage=row_Z,
                         col_linkage=col_Z,
                         figsize=figsize,
                         colors_ratio=colors_ratio,
                         dendrogram_ratio=dendrogram_ratio,
                         row_colors=row_colors,
                         col_colors=col_colors,
                         center=center,
                         cmap=cmap,
                         cbar_pos=cbar_pos,
                         **kwargs
    )
    fig.ax_heatmap.tick_params(axis='x', which='both', labelbottom=xticklabels)
    fig.ax_heatmap.tick_params(axis='y', which='both', labelbottom=yticklabels)

    # # Add a tick to every row
    # fig.ax_heatmap.set_yticks(np.arange(min(fig.ax_heatmap.get_ylim()) + 0.5, max(fig.ax_heatmap.get_ylim()) + 0.5, 1))

    if show:
        if verbose:
            utils.tprint('Showing Figure')
        plt.show()
    else:
        if verbose:
            utils.tprint('Done')

    if to_rearrange_rows:
        if verbose:
            utils.tprint('Rearranging rows')
        data, _, _ = rearrange_by_Z(data, data_arr, rownames, row_Z, row=True)
    if to_rearrange_cols:
        if verbose:
            utils.tprint('Rearranging cols')
        data, _, _ = rearrange_by_Z(data, data_arr, colnames, col_Z, row=False)

    return data, fig, row_Z, col_Z


count = int(input())
for n in range(count):
    info = input().split()
    n1 = int(info[0])
    n2 = int(info[1])
    a1 = [int(x) for x in input().split()]
    a2 = [int(x) for x in input().split()]
    a2.sort()
    i, j = 0, 0
    last = n1 - 1
    while i < n1 and j < n2 and last >= 0:
        if a1[i] > a2[j]:
            a1[last], a2[j] = a2[j], a1[last]
            last -= 1
            j += 1
        else:
            i += 1
        if last < i:
            break
    a1.sort()
    a2.sort()
    print(*a1, end=' ')
    print(*a2,end=' \n')


import sys
import re

def is_sum(xs, number):
    xs = sorted(xs)
    for ix, n in enumerate(xs[:-1]):
        for m in xs[ix+1:]:
            if n+m > number:
                break
            if n+m == number:
                return True
    return False

def find_contiguous(xs, number):
    for ix, n in enumerate(xs[:-1]):
        acc = [n]
        for m in xs[ix+1:]:
            acc.append(m)
            current_sum = sum(acc)
            if current_sum == number:
                acc = sorted(acc)
                return acc[0]+acc[-1]
            if current_sum > number:
                break

with open('input.txt') as f:
    numbers = [int(l) for l in f.readlines()]

prefix_length = 25
for current_pos, number in enumerate(numbers[prefix_length:]):
    prefix = numbers[current_pos:(prefix_length+current_pos)]

    if not is_sum(prefix, number):
        print(f'Challenge 1: {number}')
        print(f'Challenge 2: {find_contiguous(numbers, number)}')
        break


import torch

class FGM():
    def __init__(self, model):
        self.model = model
        self.backup = {}

    def attack(self, epsilon=1., emb_name='bert.embeddings.word_embeddings'):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = epsilon * param.grad / norm
                    param.data.add_(r_at)

    def restore(self, emb_name='char_emb.'):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}

class PGD():
    def __init__(self, model):
        self.model = model
        self.emb_backup = {}
        self.grad_backup = {}

    def attack(self, epsilon=1., alpha=0.3, emb_name='char_emb.', is_first_attack=False):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                if is_first_attack:
                    self.emb_backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = alpha * param.grad / norm
                    param.data.add_(r_at)
                    param.data = self.project(name, param.data, epsilon)

    def restore(self, emb_name='char_emb.'):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                assert name in self.emb_backup
                param.data = self.emb_backup[name]
        self.emb_backup = {}

    def project(self, param_name, param_data, epsilon):
        r = param_data - self.emb_backup[param_name]
        if torch.norm(r) > epsilon:
            r = epsilon * r / torch.norm(r)
        return self.emb_backup[param_name] + r

    def backup_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.grad_backup[name] = param.grad.clone()

    def restore_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.grad = self.grad_backup[name]


from slippi import Game
from time import mktime
from os import listdir, rename
from unicodedata import normalize

def parseGame(fileName):
    data = Game(fileName)

    #Haven't implemented teams file renaming yet, unsure what the format would look like
    if not data.start.is_teams:
        characters = []

        #Goes through every player in the list, if a player is NoneType, then no one was playing on that port
        for player in data.start.players:
            if player is None:
                continue

            #Saves the character and tag of the player, and makes sure that the tag is in a readable form
            characterID = player.character

            #Supports the latin characters in the name select
            characterTag = normalize('NFKC', player.tag)

            #Adds the prefix if the players tag exists, otherwise does not list the tag
            if characterTag != '':
                characterTag = '_' + characterTag

            #Adds to the list of players
            characters+=[f'{characterID.name}{characterTag}']

        #Joins the players into a single string
        charString = 'vs'.join(characters)
        stageString = data.start.stage.name

        #Only necessary to avoid duplicate file names
        timeString = int(mktime(data.metadata.date.timetuple()))

        #Constructs the new file name
        #Time string comes first in order to preserve the order of the replays
        return f'{timeString}_{charString}_{stageString}.slp'
    return None

#Goes through every file in the working directory with the 'slp' file extension and renames them
for file in listdir('.'):
    if file.endswith(".slp"):
        rename(file,parseGame(file))


"""Patchable elements module."""
from typing import Optional
import pandas as _pd
from .. import ureg as _ureg
from .. import Q_ as _Q
from georges_core.frame import Frame as _Frame

class Patchable:
    """Patchable elements are Zgoubi commands that affect the placement of the reference frame.

    A default implementation of the placement methods is provided for subclasses. It only places the entrance frame
    at the location of the placement frame and all other frames are set to the entrance frame ('point-like' element).
    """

    def __init__(self):
        """Initializes a un-patched patchable element."""
        self.LABEL1 = None
        self._entry: Optional[_Frame] = None
        self._entry_patched: Optional[_Frame] = None
        self._exit: Optional[_Frame] = None
        self._exit_patched: Optional[_Frame] = None
        self._frenet: Optional[_Frame] = None
        self._center: Optional[_Frame] = None
        self._reference_trajectory: Optional[_pd.DataFrame] = None

    def adjust_tracks_variables(self, tracks: _pd.DataFrame):
        t = tracks[tracks.LABEL1 == self.LABEL1]
        tracks.loc[tracks.LABEL1 == self.LABEL1, 'SREF'] = t['X'] + self.entry_s.m_as('m')
        tracks.loc[tracks.LABEL1 == self.LABEL1, 'YT'] = t['Y']
        try:
            tracks.loc[tracks.LABEL1 == self.LABEL1, 'YT0'] = t['Yo']
        except KeyError:
            pass
        tracks.loc[tracks.LABEL1 == self.LABEL1, 'ZT'] = t['Z']
        try:
            tracks.loc[tracks.LABEL1 == self.LABEL1, 'ZT0'] = t['Zo']
        except KeyError:
            pass

    def place(self, frame: _Frame):
        """Place the element with a reference frame.

        All the frames of the element are reset and the entrance frame is then placed with respect to the reference
        frame.

        Args:
            frame: the reference frame for the placement of the entrance frame.
        """
        self.clear_placement()
        self._entry = frame.__class__(frame)

    def clear_placement(self):
        """Clears all the frames."""
        self._entry = None
        self._entry_patched = None
        self._exit = None
        self._exit_patched = None
        self._center = None

    @property
    def length(self) -> _Q:
        """Length of the element.

        Returns:
            the length of the element with units.
        """
        return 0.0 * _ureg.cm

    @property
    def entry(self) -> Optional[_Frame]:
        """Entrance frame.

        Returns:
            the frame of the entrance of the element.
        """
        return self._entry

    @property
    def entry_patched(self) -> Optional[_Frame]:
        """Entrance patched frame.

        Returns:
            the frame of the entrance of the element with the patch applied.
        """
        if self._entry_patched is None:
            self._entry_patched = self.entry.__class__(self.entry)
        return self._entry_patched

    @property
    def exit(self) -> Optional[_Frame]:
        """Exit frame.

        Returns:
            the frame of the exit of the element.
        """
        if self._exit is None:
            self._exit = self.entry_patched.__class__(self.entry_patched)
        return self._exit

    @property
    def exit_patched(self) -> Optional[_Frame]:
        """Exit patched frame.

        Returns:
            the frame of the exit of the element with the patch applied.
        """
        if self._exit_patched is None:
            self._exit_patched = self.exit.__class__(self.exit)
        return self._exit_patched

    @property
    def frenet_orientation(self) -> Optional[_Frame]:
        if self._frenet is None:
            self._frenet = self.entry_patched.__class__(self.entry_patched)
        return self._frenet

    @property
    def center(self) -> Optional[_Frame]:
        """Center frame.

        Returns:
            the frame of the center of the element.
        """
        if self._center is None:
            self._center = self.entry.__class__(self.entry)
        return self._center

    @property
    def reference_trajectory(self) -> _pd.DataFrame:
        """

        Returns:

        """
        return self._reference_trajectory

    @reference_trajectory.setter
    def reference_trajectory(self, ref: _pd.DataFrame):
        """

        Args:
            ref:

        Returns:

        """
        self._reference_trajectory = ref

    @property
    def entry_s(self) -> Optional[_ureg.Quantity]:
        """

        Returns:

        """
        if self.reference_trajectory is not None:
            return self.reference_trajectory['S'].min() * _ureg.m
        else:
            return 0.0 * _ureg.m

    @property
    def exit_s(self) -> Optional[_ureg.Quantity]:
        """

        Returns:

        """
        if self.reference_trajectory is not None:
            return self.reference_trajectory['S'].max() * _ureg.m
        else:
            return 0.0 * _ureg.m

    @property
    def optical_length(self) -> Optional[_ureg.Quantity]:
        """

        Returns:

        """
        if self.reference_trajectory is not None:
            return (self.reference_trajectory['S'].max() - self.reference_trajectory['S'].min()) * _ureg.m
        else:
            return 0.0 * _ureg.m


import unittest
import warnings
import datetime

from unittest.mock import patch, MagicMock
from io import BytesIO

from control.account import (
    DropboxAccountWrapper, GoogleDriveAccountWrapper,
    InterruptibleGoogleDriveUploadFileHandle, InterruptibleGoogleDriveDownloadFileHandle,
    TaskInterruptedException, CloudAPIFactory
)
from model.account import AccountTypes, AccountData
from model.task import Task
from model.message import MessageTypes
from model.file import FilePart

from dropbox.files import FileMetadata, FolderMetadata
from Crypto.Cipher import AES

warnings.filterwarnings("ignore", category=DeprecationWarning)

class TestDropboxAccountWrapper(unittest.TestCase):

    def setUp(self):
        self.testAccountData = AccountData(
            AccountTypes.Dropbox,
            "testIdentifier",
            "sixteen byte key",
            {"apiToken": "testApiToken"},
            1
        )

    @patch("dropbox.Dropbox")
    def test_getFileList_returns_encrypted_files_recursively(self, mockDropbox):
        testCloudFiles = [
            FileMetadata(name="apple.txt__1__1.enc", size=5, client_modified=datetime.datetime(2020, 1, 1, 6, 10, 15), path_display="/apple.txt__1__1.enc"),
            FolderMetadata(name="thisFolderShouldBeIgnored"),
            FileMetadata(name="pear.txt__1__2.enc", size=10, client_modified=datetime.datetime(2020, 1, 1, 7, 10, 15), path_display="/subDir/pear.txt__1__2.enc"),
            FileMetadata(name="notAnEncodedFile.jpeg", size=20, client_modified=datetime.datetime(2020, 1, 1, 7, 10, 15), path_display="/notAnEncodedFile.jpeg"),
        ]
        mockDropbox.return_value.files_list_folder.return_value.entries = testCloudFiles
        cloudAccount = DropboxAccountWrapper(self.testAccountData)
        files = cloudAccount.getFileList()

        self.assertEqual(len(files), 2)
        self.assertEqual(mockDropbox.return_value.files_list_folder.call_args[0][0], "")
        self.assertEqual(mockDropbox.return_value.files_list_folder.call_args[1], {"recursive": True})

        self.assertEqual(files[0].filename, testCloudFiles[0].name)
        self.assertEqual(files[0].modified, int(testCloudFiles[0].client_modified.timestamp()))
        self.assertEqual(files[0].size, testCloudFiles[0].size)
        self.assertEqual(files[0].path, "")
        self.assertEqual(files[0].fullPath, "apple.txt__1__1.enc")
        self.assertEqual(files[0].storingAccountID, self.testAccountData.id)

        self.assertEqual(files[1].filename, testCloudFiles[2].name)
        self.assertEqual(files[1].modified, int(testCloudFiles[2].client_modified.timestamp()))
        self.assertEqual(files[1].size, testCloudFiles[2].size)
        self.assertEqual(files[1].path, "subDir")
        self.assertEqual(files[1].fullPath, "subDir/pear.txt__1__2.enc")
        self.assertEqual(files[1].storingAccountID, self.testAccountData.id)

    @patch("dropbox.files.UploadSessionCursor")
    @patch("dropbox.Dropbox")
    def test_upload_writes_iv_and_encrypts_file_data(self, mockDropbox, mockUploadCursor):
        uploadSessionStarter = MagicMock()
        uploadSessionAppender = MagicMock()

        mockDropbox.return_value.files_upload_session_start = uploadSessionStarter
        mockDropbox.return_value.files_upload_session_finish = uploadSessionAppender

        testSecretData = b"secret test data"
        testFileHandle = BytesIO(testSecretData)
        testPartName = "testFile__1__1.enc"
        testClientModified = datetime.datetime(2020, 1, 1, 10, 5, 30)

        testFileData = {"userTimezone": "+0200", "utcModified": testClientModified.timestamp(), "path": "subDir"}
        testTask = Task(taskType=MessageTypes.UPLOAD_FILE, data=testFileData)

        cloudAccount = DropboxAccountWrapper(self.testAccountData)
        result = cloudAccount.upload(testFileHandle, len(testSecretData), testPartName, testTask)

        cipher = AES.new(self.testAccountData.cryptoKey.encode(), AES.MODE_CFB, iv=uploadSessionStarter.call_args[0][0])
        encoded = mockDropbox.return_value.files_upload_session_finish.call_args[0][0]

        self.assertNotEqual(encoded, testSecretData)
        self.assertEqual(cipher.decrypt(encoded), testSecretData)

    @patch("dropbox.files.UploadSessionCursor")
    @patch("dropbox.Dropbox")
    def test_upload_does_nothing_if_task_is_stale(self, mockDropbox, mockUploadCursor):
        uploadSessionStarter = MagicMock()
        uploadSessionAppender = MagicMock()

        mockDropbox.return_value.files_upload_session_start = uploadSessionStarter
        mockDropbox.return_value.files_upload_session_finish = uploadSessionAppender

        testSecretData = b"secret test data"
        testFileHandle = BytesIO(testSecretData)
        testPartName = "testFile__1__1.enc"
        testClientModified = datetime.datetime(2020, 1, 1, 10, 5, 30)

        testFileData = {"userTimezone": "+0200", "utcModified": testClientModified.timestamp(), "path": "subDir"}
        testTask = Task(taskType=MessageTypes.UPLOAD_FILE, data=testFileData, stale=True)

        cloudAccount = DropboxAccountWrapper(self.testAccountData)
        result = cloudAccount.upload(testFileHandle, len(testSecretData), testPartName, testTask)

        self.assertEqual(uploadSessionStarter.call_count, 0)
        self.assertEqual(uploadSessionAppender.call_count, 0)
        self.assertIsNone(result)

    @patch("requests.get")
    def test_download_file_decodes_content_if_not_interrupted(self, mockRequest):
        secretData = b"secret test data"
        testEncoder = AES.new(self.testAccountData.cryptoKey.encode(), AES.MODE_CFB)
        testIV = testEncoder.iv

        ivResponse = MagicMock()
        ivResponse.content = testIV

        encryptedDataResponse = MagicMock()
        encryptedDataResponse.content = testEncoder.encrypt(secretData)

        mockRequest.side_effect = [ivResponse, encryptedDataResponse]

        testDownloadFileHandle = BytesIO()
        testFilePartInfo = FilePart(
            filename="apple.txt__1__1.enc", modified=int(datetime.datetime(2020, 1, 1, 10, 10, 30).timestamp()),
            size=len(encryptedDataResponse.content) + len(testIV), path="", fullPath="apple.txt__1__1.enc",
            storingAccountID=self.testAccountData.id, extraInfo={}
        )

        testDownloadFileTask = Task(taskType=MessageTypes.DOWNLOAD_FILE)
        cloudAccount = DropboxAccountWrapper(self.testAccountData)
        cloudAccount.download(testDownloadFileHandle, testFilePartInfo, testDownloadFileTask)

        testDownloadFileHandle.seek(0)

        self.assertEqual(secretData, testDownloadFileHandle.read())

    @patch("requests.get")
    def test_download_file_does_nothing_if_task_is_stale(self, mockRequest):
        secretData = b"secret test data"
        testEncoder = AES.new(self.testAccountData.cryptoKey.encode(), AES.MODE_CFB)
        testIV = testEncoder.iv

        ivResponse = MagicMock()
        ivResponse.content = testIV

        encryptedDataResponse = MagicMock()
        encryptedDataResponse.content = testEncoder.encrypt(secretData)

        mockRequest.side_effect = [ivResponse, encryptedDataResponse]

        testDownloadFileHandle = BytesIO()
        testFilePartInfo = FilePart(
            filename="apple.txt__1__1.enc", modified=int(datetime.datetime(2020, 1, 1, 10, 10, 30).timestamp()),
            size=len(encryptedDataResponse.content) + len(testIV), path="", fullPath="apple.txt__1__1.enc",
            storingAccountID=self.testAccountData.id, extraInfo={}
        )

        testDownloadFileTask = Task(taskType=MessageTypes.DOWNLOAD_FILE, stale=True)
        cloudAccount = DropboxAccountWrapper(self.testAccountData)
        cloudAccount.download(testDownloadFileHandle, testFilePartInfo, testDownloadFileTask)

        testDownloadFileHandle.seek(0)

        self.assertEqual(b"", testDownloadFileHandle.read())

    @patch("dropbox.Dropbox")
    def test_deleteFile_delegates_call_properly(self, mockDropbox):

        mockDropbox.return_value.files_delete = MagicMock()
        testFilePartInfo = FilePart(
            filename="apple.txt__1__1.enc", modified=int(datetime.datetime(2020, 1, 1, 10, 10, 30).timestamp()),
            size=32, path="subDir", fullPath="subDir/apple.txt__1__1.enc",
            storingAccountID=self.testAccountData.id, extraInfo={}
        )

        cloudAccount = DropboxAccountWrapper(self.testAccountData)
        cloudAccount.deleteFile(testFilePartInfo)

        self.assertEqual(mockDropbox.return_value.files_delete.call_count, 1)
        self.assertEqual(mockDropbox.return_value.files_delete.call_args[0][0], f"/{testFilePartInfo.fullPath}")

    @patch("dropbox.Dropbox")
    def test_moveFile_delegates_call_properly(self, mockDropbox):
        mockDropbox.return_value.files_move = MagicMock()
        testFilePartInfo = FilePart(
            filename="apple.txt__1__1.enc", modified=int(datetime.datetime(2020, 1, 1, 10, 10, 30).timestamp()),
            size=32, path="", fullPath="apple.txt__1__1.enc",
            storingAccountID=self.testAccountData.id, extraInfo={}
        )
        destinationPath = "subDir/apple.txt__1__2.enc"

        cloudAccount = DropboxAccountWrapper(self.testAccountData)
        cloudAccount.moveFile(testFilePartInfo, destinationPath)

        self.assertEqual(mockDropbox.return_value.files_move.call_count, 1)
        self.assertEqual(mockDropbox.return_value.files_move.call_args[0][0], f"/{testFilePartInfo.fullPath}")
        self.assertEqual(mockDropbox.return_value.files_move.call_args[0][1], f"/{destinationPath}")

class TestInterruptibleGoogleDriveFileHandles(unittest.TestCase):

    def test_upload_fileHandle_raises_exception_if_task_is_stale(self):
        fakeFileHandle = BytesIO(b"lorem ipsum")
        testFileData = {"userTimezone": "+0200", "utcModified": datetime.datetime(2020, 1, 1, 10, 5, 30).timestamp(), "path": "subDir"}
        testTask = Task(taskType=MessageTypes.UPLOAD_FILE, data=testFileData, stale=True)

        testInterruptibleHandle = InterruptibleGoogleDriveUploadFileHandle(fakeFileHandle, testTask)
        try:
            testInterruptibleHandle.read(8192)
            self.fail("InterruptibleGoogleDriveUploadFileHandle failed to raise TaskInterruptedException on a stale task!")
        except TaskInterruptedException:
            pass

    def test_download_fileHandle_raises_exception_if_task_is_stale(self):
        fakeDownloadFileHandle = BytesIO()
        testDownloadFileTask = Task(taskType=MessageTypes.DOWNLOAD_FILE, stale=True)

        testInterruptibleHandle = InterruptibleGoogleDriveDownloadFileHandle(fakeDownloadFileHandle, "sixteen byte key", testDownloadFileTask)
        try:
            testInterruptibleHandle.write(b"fakeinitvectorrr")
            testInterruptibleHandle.write(b"fake file data should raise exception")
            self.fail("InterruptibleGoogleDriveDownloadFileHandle failed to raise TaskInterruptedException on a stale task!")
        except TaskInterruptedException:
            pass

if __name__ == '__main__':
    unittest.main()


from slugify import slugify
from sqlalchemy import case, desc, func, or_
from sqlalchemy.orm import Session, Query

from models.game import Game
from models.user import User
from schemas.user import UserCreate, UserProfile, UserStat

def list_users(db: Session):
    return db.query(User).order_by(User.display_name).all()

def create_user(db: Session, user: UserCreate) -> User:
    db_user = User(
        display_name=user.display_name,
        google_id=user.google_id,
        role=user.role,
        slug_name=slugify(user.display_name)
    )
    db.add(db_user)
    db.commit()
    db.refresh(db_user)
    return db_user

def update_user(db: Session, user: User) -> User:
    db.merge(user)
    db.commit()
    db.refresh(user)
    return user

def get_user_from_id(db: Session, user_id: int) -> User:
    return db.get(User, user_id)

def get_user_from_google_id(db: Session, google_id: str) -> User:
    return db.query(User).filter(User.google_id == google_id).first()

def get_user_from_slug_name(db: Session, slug_name: str) -> User:
    return db.query(User).filter(User.slug_name == slug_name).first()

def get_stats_for_user(db: Session, user_id: int):
    user = get_user_from_id(db, user_id)
    if not user:
        return None
    teammate_query = func.sum(
        case(
            (user_id == Game.player_a1_id, case((Game.player_a2_id == User.id, 1), else_=0)),
            (user_id == Game.player_a2_id, case((Game.player_a1_id == User.id, 1), else_=0)),
            (user_id == Game.player_b1_id, case((Game.player_b2_id == User.id, 1), else_=0)),
            (user_id == Game.player_b2_id, case((Game.player_b1_id == User.id, 1), else_=0))
        )
    ).label("teammate")
    opponent_query = func.sum(
        case(
            (user_id == Game.player_a1_id, case((Game.player_a2_id != User.id, 1), else_=0)),
            (user_id == Game.player_a2_id, case((Game.player_a1_id != User.id, 1), else_=0)),
            (user_id == Game.player_b1_id, case((Game.player_b2_id != User.id, 1), else_=0)),
            (user_id == Game.player_b2_id, case((Game.player_b1_id != User.id, 1), else_=0))
        )
    ).label("opponent")
    wins_query = func.sum(
        case(
            (user_id == Game.player_a1_id,
             case((Game.player_a2_id != User.id, case((Game.a_won, 1), else_=0)), else_=0)),
            (user_id == Game.player_a2_id,
             case((Game.player_a1_id != User.id, case((Game.a_won, 1), else_=0)), else_=0)),
            (user_id == Game.player_b1_id,
             case((Game.player_b2_id != User.id, case((Game.b_won, 1), else_=0)), else_=0)),
            (
                user_id == Game.player_b2_id,
                case((Game.player_b1_id != User.id, case((Game.b_won, 1), else_=0)), else_=0))
        )
    ).label("wins")
    query = (db.query(
        User.id,
        User.display_name,
        teammate_query,
        opponent_query,
        wins_query,
        (opponent_query - wins_query).label("defeats")
    )
     .filter(User.id != user_id)
     .filter(or_(Game.player_a1_id == user_id,
                 Game.player_a2_id == user_id,
                 Game.player_b1_id == user_id,
                 Game.player_b2_id == user_id,
                 ))
     .filter(or_(
                Game.player_a1_id == User.id,
                Game.player_a2_id == User.id,
                Game.player_b1_id == User.id,
                Game.player_b2_id == User.id,
                ))
     .group_by(User.id)
     .group_by(User.display_name)
     )
    teammate = query.order_by(teammate_query.desc()).first()
    opponent = query.order_by(opponent_query.desc()).first()
    target = query.order_by(wins_query.desc()).first()
    executioner = query.order_by(desc("defeats")).first()
    simple_stats = get_simple_stat_query(db).filter(User.id == user_id).first()
    if not simple_stats:
        return UserProfile(
            id=user.id, display_name=user.display_name
        )
    return UserProfile(
        id=simple_stats.id,
        display_name=simple_stats.display_name,
        wins=simple_stats.wins,
        games=simple_stats.games,
        stars=simple_stats.stars,
        best_teammate=teammate.display_name,
        teammate_times=teammate.teammate,
        best_opponent=opponent.display_name,
        oppositions=opponent.opponent,
        best_target=target.display_name,
        victories=target.wins,
        executioner=executioner.display_name,
        defeats=executioner.defeats
    )

def get_stats_for_users(db: Session) -> list[UserStat]:
    query = get_simple_stat_query(db)
    users = query.all()
    response = []
    for user in users:
        response.append(
            UserStat(id=user.id, display_name=user.display_name, wins=user.wins, games=user.games, stars=user.stars))
    return response

def get_simple_stat_query(db: Session) -> Query:
    wins = func.sum(
        case(
            (User.id == Game.player_a1_id, case((Game.a_won, 1), else_=0)),
            (User.id == Game.player_a2_id, case((Game.a_won, 1), else_=0)),
            else_=case((Game.b_won, 1), else_=0)
        )
    ).label("wins")
    stars = func.sum(
        case(
            (User.id == Game.player_a1_id, Game.stars_a),
            (User.id == Game.player_a2_id, Game.stars_a),
            else_=Game.stars_b
        )
    ).label("stars")
    query = (db.query(
        User.id,
        User.display_name,
        func.count(Game.id).label("games"),
        wins,
        stars
    ).filter(or_(
        Game.player_a1_id == User.id,
        Game.player_a2_id == User.id,
        Game.player_b1_id == User.id,
        Game.player_b2_id == User.id,
    )).group_by(User.id)
             .group_by(User.display_name)
             .order_by(wins.desc())
             )
    return query

def delete_user(db: Session, user: User):
    """
    Works only if the user has not participated in any game
    """
    db.delete(user)
    db.commit()


#
# Created by Maksim Eremeev (mae9785@nyu.edu)
#

from ..elsa import Preprocessing
from tqdm.auto import tqdm
import argparse

def prepare_fasttext_cnn(preprocessor, dataset_path, save_path):
    with open(dataset_path, 'r') as f:
        for line in tqdm(f, total=221000):
            doc = preprocessor.preproc(line, check_stopwords=False, check_length=True,
                                       use_lemm=False, use_stem=False, include_tf=False)
            with open(save_path, 'a') as fw:
                fw.write(f'{doc}\n')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--stopwords', nargs='*', help='path to stopwords.txt')
    parser.add_argument('-d', '--dataset', nargs='*', help='path to the dataset')
    parser.add_argument('-o', '--output', nargs='*', help='output path')
    args = parser.parse_args()

    preprocessor = Preprocessing(stopwords=args.stopwords[0])
    dataset_path = args.dataset[0]
    save_path = args.output[0]

    prepare_fasttext_cnn(preprocessor, dataset_path, save_path)


import numpy as np
import os
from argparse import ArgumentParser
from neural_net.nn_utils import get_trained_network

from utils.delegates_ls import *
from sampling.rate_schedulers import *
from models.model_generators import *

# ranges for each type of tuning
ranges = {
    "n_iters" : np.linspace(1000, 10000, num=10),
    "rate" : [8e-7],
    "batch_size" : np.linspace(1000, 8000, num=8),
    "nesterov" : [0.5, 0.75, 0.9, 0.99]
}

def main(param, output_path, config_file, problem, debug_steps):
    # read from config
    config = configparser.ConfigParser()
    config.read(config_file)
    batch_size = config.getint('sampling.least-squares', 'batch_size')
    n_iters = config.getint('sampling.least-squares', 'n_iters')
    scheduler = scheduler_from_string(config.get('sampling.least-squares', 'scheduler'))
    rate = config.getfloat('sampling.least-squares', 'rate')
    momentum_type = config.get('sampling.least-squares', 'momentum')
    beta = config.getfloat('sampling.least-squares', 'beta')

    # warning for momentum type
    if momentum_type != 'nesterov':
        print(f"WARN: Nesterov momentum is used for least-squares tuning, overriding '{momentum_type}' specified in config")

    # set up params
    params = {
        "n_iters" : n_iters,
        "scheduler" : scheduler,
        "rate" : rate,
        "batch_size" : batch_size,
        "nesterov" : beta,
        "scheduler_n_epochs" : 10000
    }

    # set up models and sampler
    print("INFO: Training Network")
    if problem == 'nn':
        net = get_trained_network(config_file)
        model, _ = generate_nn_linear_models(config_file, net)
    elif problem == 'synthetic':
        model, _ = generate_synthetic_models(config_file)

    # sample for parameter range
    inputs = ranges[param]
    print(inputs)
    errors = np.zeros((len(inputs), n_iters))
    if param == 'batch_size':
        # compute samples and errors
        for i, p in enumerate(inputs):
            print("Computing for batch size = {}".format(p))
            _, debug = batch_size_delegate(model, p, debug_steps=debug_steps, params=params)
            errors[i] = debug[1]
            # make new folder for outputs
            path = f"{output_path}/{param}"
            if not os.path.exists(path):
                raise Exception(f"{path} directory required, please make it and re-run")
            output = params.copy()
            output[param] = p
            output['errors'] = errors[i]
            output['scheduler'] = config.get('sampling.least-squares', 'scheduler') # get scheduler as string
            df = pd.DataFrame(output)
            df.to_csv(f'{path}/ls_tuning_{param}_{i}.csv',index=False)
    elif param == 'rate':
        # compute samples and errors
        for i, p in enumerate(inputs):
            print("Computing for learning rate = {}".format(p))
            print("INFO: Generating Linear Model")
            _, debug = learning_rate_delegate(model, p, debug_steps=debug_steps, params=params)
            errors[i] = debug[1]
            # make new folder for outputs
            path = f"{output_path}/{param}"
            if not os.path.exists(path):
                raise Exception(f"{path} directory required, please make it and re-run")
            output = params.copy()
            output[param] = p
            output['errors'] = errors[i]
            output['scheduler'] = config.get('sampling.least-squares', 'scheduler') # get scheduler as string
            df = pd.DataFrame(output)
            df.to_csv(f'{path}/ls_tuning_{param}_{batch_size}_{i}.csv',index=False)
    elif param == 'nesterov':
        # compute samples and errors
        for i, p in enumerate(inputs):
            print("Computing for nesterov parameter = {}".format(p))
            print("INFO: Generating Linear Model")
            _, debug = nesterov_momentum_delegate(model, p, debug_steps=debug_steps, params=params)
            errors[i] = debug[1]
            # make new folder for outputs
            path = f"{output_path}/{param}"
            if not os.path.exists(path):
                raise Exception(f"{path} directory required, please make it and re-run")
            output = params.copy()
            output[param] = p
            output['errors'] = errors[i]
            output['scheduler'] = config.get('sampling.least-squares', 'scheduler') # get scheduler as string
            df = pd.DataFrame(output)
            df.to_csv(f'{path}/ls_tuning_{param}_{i}.csv',index=False)
    else:
        raise Exception("Unknown parameter specified")

    print("Successfully written output to {}".format(path))

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--param', type=str)
    parser.add_argument('--problem', type=str, default='nn')
    parser.add_argument('--output_path', type=str, default='ls_outputs')
    parser.add_argument('--config_file', type=str, default='configs/nn_config.ini')
    parser.add_argument('--sampler_debug_steps', type=int, default=1000)
    args = parser.parse_args()
    main(args.param, args.output_path, args.config_file, args.problem, args.sampler_debug_steps)


import tensorflow as TF
import os
import random

from spec_augment import initialize_augmentation

# -------------------- CONSTANTS -------------------- #

# TFRecord Directories
RECORD_DIR = os.path.join('D:', 'Track_Records')
TRAIN_DIR = os.path.join(RECORD_DIR, 'train')
TEST_DIR = os.path.join(RECORD_DIR, 'test')

# Training & Testing Dataset Sizes
N_TRAIN = 71_421
N_TEST = 7_714
N_RECORD = 1_000

# Spectrogram Constants
DIM = (430, 84)
SPEC_AUGMENT = None

# TFRecord Byte Formatting
FEATURE_STRUCTURE = {
    'parameters': TF.io.FixedLenFeature([DIM[0] * DIM[1]], TF.float32),
    'label': TF.io.FixedLenFeature([], TF.int64)
}

# -------------------- MAIN FUNCTION --------------------- #

def train_model(model, optimizer = TF.keras.optimizers.Adam(lr = 1e-3),
                batch_size = 64, epochs = 10, callbacks = None, augment = True):
    # Augmentation Initialization
    if augment:
        SPEC_AUGMENT = initialize_augmentation(DIM)

    # Dataset Generation
    TRAIN_data = generate_TRAIN(batch_size, epochs)
    TEST_data = generate_TEST(batch_size, epochs)

    # Compilation
    model.compile(
            optimizer,
            loss = 'sparse_categorical_crossentropy',
            metrics = ['accuracy']
    )

    # Training
    history = model.fit(
        TRAIN_data,
        validation_data = TEST_data,
        epochs = epochs,
        steps_per_epoch = (N_TRAIN // batch_size),
        validation_steps = (N_TEST // batch_size),
        verbose = 1,
        callbacks = callbacks
    )

    # Returning the training results.
    return history.history

# ----------------- DATASET GENERATION ----------------- #

def generate_TRAIN(batch_size, epochs):
    files = [os.path.join(TRAIN_DIR, file) for file in os.listdir(TRAIN_DIR)]

    # Sharded & Shuffled Dataset
    dataset = TF.data.Dataset.from_tensor_slices(files)
    dataset = dataset.shuffle(len(files))

    # Interleaved, Repeated, & Shuffled Dataset
    dataset = TF_interleave(dataset, file_dataset)
    dataset = dataset.repeat(epochs)
    dataset = dataset.shuffle(N_RECORD)

    # Parsed & Augmented Dataset
    dataset = TF_map(dataset, parse_dataset)
    dataset = TF_map(dataset, augment)
    dataset = dataset.apply(TF.data.experimental.ignore_errors())

    # Batched Dataset
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(TF.data.experimental.AUTOTUNE)

    return dataset

def generate_TEST(batch_size, epochs):
    files = [os.path.join(TEST_DIR, file) for file in os.listdir(TEST_DIR)]

    # Interleaved & Repeated Dataset
    dataset = TF.data.Dataset.from_tensor_slices(files)
    dataset = TF_interleave(dataset, file_dataset)
    dataset = dataset.repeat(epochs)

    # Parsed Dataset
    dataset = TF_map(dataset, parse_dataset)

    # Batched Dataset
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(TF.data.experimental.AUTOTUNE)

    return dataset

# ----------------- FUNCTION MAPPERS ----------------- #

def TF_interleave(dataset, map_func):
    return dataset.interleave(map_func,
              num_parallel_calls = TF.data.experimental.AUTOTUNE, deterministic = False)

def TF_map(dataset, map_func):
    return dataset.map(map_func,
               num_parallel_calls = TF.data.experimental.AUTOTUNE, deterministic = False)

# ----------------- MAPPING FUNCTIONS ----------------- #

#
def file_dataset(file):
    return TF.data.TFRecordDataset(file)

def parse_dataset(feature):
    data = TF.io.parse_single_example(feature, FEATURE_STRUCTURE)
    data['parameters'] = TF.reshape(data['parameters'], DIM)
    return data['parameters'], data['label']

def augment(image, label):
    if SPEC_AUGMENT is not None:
        image = SPEC_AUGMENT(image)
    return image, label


from __future__ import print_function

import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML
import torchsummary

# Set random seem for reproducibility
manualSeed = 999
#manualSeed = random.randint(1, 10000) # use if you want new results
random.seed(manualSeed)
torch.manual_seed(manualSeed)

# Root directory for dataset
dataroot = "split/DSMWRAPPER"

# Number of workers for dataloader
workers = 2

# Batch size during training
batch_size = 128

# Spatial size of training images. All images will be resized to this
#   size using a transformer.
image_size = 64

# Number of channels in the training images. For color images this is 3
nc = 1

# Size of z latent vector (i.e. size of generator input)
nz = 100

# Size of feature maps in generator
ngf = 64

# Size of feature maps in discriminator
ndf = 64

# Number of training epochs
num_epochs = 5

# Learning rate for optimizers
lr = 0.0002

# Beta1 hyperparam for Adam optimizers
beta1 = 0.5

# Number of GPUs available. Use 0 for CPU mode.
ngpu = 1

# We can use an image folder dataset the way we have it setup.
# Create the dataset
dataset = dset.ImageFolder(root=dataroot,
                           transform=transforms.Compose([
                               transforms.Grayscale(num_output_channels=1),
                               transforms.Resize(image_size),
                               transforms.CenterCrop(image_size),
                               transforms.ToTensor(),
                               #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                           ]))

# Create the dataloader
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                         shuffle=True, num_workers=workers)

# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")

# Plot some training images
real_batch = next(iter(dataloader))
plt.figure(figsize=(8,8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))

netG = Generator(1).to(device)

netG.apply(weights_init)

torchsummary.summary(netG,input_size=(100,1,1))

netD = Discriminator(ngpu).to(device)

netD.apply(weights_init)

torchsummary.summary(netD,input_size=(1,64,64))

# Initialize BCELoss function
criterion = nn.BCELoss()

# Create a batch of latent vectors that we will use to visuallize
# as we train the generator
fixed_noise = torch.randn(64, nz, 1, 1, device=device)

# Set labelling convention
real_label = 1
fake_label = 0

# Set up the two optimizers for G and D, we use ADAM since it's in the paper
optimizerD = optim.Adam(netD.parameters(),lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(),lr=lr, betas=(beta1, 0.999))

# Training Loop

# lists to keep track of losses and progress
img_list = []
G_losses = []
D_losses = []
iters = 0

print("Starting Training Loop")

#for epoch in range(25):

for epoch in range(num_epochs):

    # for each batch in the dataloader
    for i, data in enumerate(dataloader,0):

        ####
        # Train the discriminator: Maximize log(D(x) + log(1- D(G(z)))
        ####

        ## Train with all real batch
        netD.zero_grad()

        # format the batch
        real_cpu = data[0].to(device)
        b_size = real_cpu.size(0) #checks to see uniform batch size
        label = torch.full((b_size,), real_label, device=device)
        # Forward pass for real batch through the Discriminator
        output = netD(real_cpu).view(-1)

        # Calculate loss
        errD_real = criterion(output, label)

        # Calculate gradients for D in backward pass
        errD_real.backward()

        D_x = output.mean().item()

        ## Train with all fake batch
        noise = torch.randn(b_size, nz, 1, 1, device=device)

        # Generate Fake images using the Generator
        fake = netG(noise)
        label.fill_(fake_label)

        # Classify fake images
        output = netD(fake.detach()).view(-1)

        # Calculate the Discriminators loss on the fake batch
        errD_fake = criterion(output,label)

        # Calculate the gradients of the backward pass for the Generator
        errD_fake.backward()
        D_G_z1 = output.mean().item()

        # Add the gradients from the real and fake batches
        errD = errD_real + errD_fake
        optimizerD.step()

        ####
        # training the generator: Maximize log(D(G(z)))
        ###

        netG.zero_grad()
        label.fill_(real_label) # fake labels are real for the generator cost

        # since we just updated D, we perform another forward pass of all fake through D
        output = netD(fake).view(-1)

        # Calculate generator loss basd on the output
        errG = criterion(output,label)

        # Calculate gradients
        errG.backward()
        D_G_z2 = output.mean().item()

        # update G
        optimizerG.step()

        # Output stats during training
        if i%50 == 0:
            print('[%d/%d][%d/%d]\tLoss_D; %.4f\tLoss_g:%.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                 % (epoch, num_epochs, i, len(dataloader),
                   errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))

        G_losses.append(errG.item())
        D_losses.append(errD.item())

        # Check how the generator is doing by saving G's output on fixed_noise
        if (iters % 500 == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)):
            with torch.no_grad():
                fake = netG(fixed_noise).detach().cpu()
            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))

        iters += 1

# Grab a batch of real images from the dataloader
real_batch = next(iter(dataloader))

# Plot the real images
plt.figure(figsize=(30,30))
plt.subplot(1,2,1)
plt.axis("off")
plt.title("Real Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(),(1,2,0)))

# Plot the fake images from the last epoch
plt.subplot(1,2,2)
plt.axis("off")
plt.title("Fake Images")
plt.imshow(np.transpose(img_list[-1],(1,2,0)))
plt.show()

torch.save(netD,"discriminator.pt")
torch.save(netG,"generator.pt")


# IN5450 Mandatory 2
# Thomas Aussaguès, 14/03/2022
# thomas.aussagues@imt-atlantique.net

# This script contains the configurations dictionaries

import numpy as np

# General parameters

config = dict()
config['DOA_array'] = np.arange(-50,50,0.01)

# Questions 1->7: incoherent sources

config1 = dict()
# Incoherent sources
config1['coherent'] = False
# Number of time samples
config1['N'] = 100
# number of sensors
config1['M'] = 10
# Source 1 SNR (in dB)
config1['SNR1'] = 0
# Source 2 SNR (in dB)
config1['SNR2'] = 0
# Source 1 direction of arrival (degrees)
config1['theta1'] = 0
# Source 2 direction of arrival (degrees)
config1['theta2'] = -10
# Phase shift between the the two sources (degrees)
config1['gamma'] = 45
# Array element spacing
config1['d'] = 0.5
# Normalized wavenumber
config1['k'] = 2*np.pi
# Time sampling interval
config1['T'] = 0.5
# Source 1 normalized frequency/pulsation
config1['omega1'] = 2*np.pi
# Here, we want to make the two sources incoherent. Therefore, we make source 2 frequency slightly different from source 1 frequency
config1['omega2'] = config1['omega1']#*(1-3/config1['N'])
# We fix the numpy seed to ensure that all sources are subject to the same noise
config1['seed'] = 0
# We generate a complex AWGN (spatially white) with variance (1 or 2)
config1['variance'] = 1

# Question 8: coherent sources

config2 = dict()
# Coherent sources
config2['coherent'] = True
# Number of time samples
config2['N'] = 100
# number of sensors
config2['M'] = 10
# Source 1 SNR (in dB)
config2['SNR1'] = 0
# Source 2 SNR (in dB)
config2['SNR2'] = 0
# Source 1 direction of arrival (degrees)
config2['theta1'] = 0
# Source 2 direction of arrival (degrees)
config2['theta2'] = 10
# Phase shift between the the two sources (degrees)
config2['gamma'] = 45
# Array element spacing
config2['d'] = 0.5
# Normalized wavenumber
config2['k'] = 2*np.pi
# Time sampling interval
config2['T'] = 0.5
# Source 1 normalized frequency/pulsation
config2['omega1'] = 2*np.pi
# Here, we want to make the two sources coherent.
config2['omega2'] = 2*np.pi
# We fix the numpy seed to ensure that all sources are subject to the same noise
config2['seed'] = 0
# We generate a complex AWGN (spatially white) with variance (1 or 2)
config2['variance'] = 1
# Diagonal loading: delta value
config2['delta'] = -15
# Spatial smoothing: number of sensors per subarray
config2['L'] = 5

def get_general_config():
    return config

def get_config_1()->dict:
    return config1

def get_config_2()->dict:
    return config2

def get_config_SNR_analysis(snr:float)->dict:
    config2['SNR1'] = snr
    config2['SNR2'] = snr
    return config2


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

import os
import requests
import time
from datetime import datetime
import pandas as pd
from bs4 import BeautifulSoup
from operator import itemgetter

# NAVITIME 住所検索TOP
def crawl():
  init()
  prefectures()

def init():
  if not os.path.exists('data'):
    os.mkdir('data')

def prefectures():
  print(f"prefectures start : {now()}")
  TOP_URL = 'https://www.navitime.co.jp/?ctl=0050'
  response = requests.get(TOP_URL)
  soup = BeautifulSoup(response.text, 'lxml')
  sub_urls = []
  prefectures = []
  for li in soup.find_all("li", attrs={"class": "pref-item"}):
    target = li.find("a", reversed=False)
    url = target.get("href")
    sub_urls.append(url)
    data = itemgetter(4, 5)(url.split("/"))
    prefectures.append(data)

  df = pd.DataFrame(prefectures, columns=['code', 'name'])
  df.to_csv("data/prefectures.csv", index=False)
  print(f"prefectures end : {now()}")

  cities(sub_urls)

def cities(urls: list):
  print(f"cities start : {now()}")
  sub_urls = []
  cities = []
  for url in urls:
    prefecture_code =  url.split("/")[4]
    print(f"----- prefecture_code={prefecture_code} start : {now()} -----")
    # if not prefecture_code == '13':
    #   continue
    response = requests.get(url)
    soup = BeautifulSoup(response.text, 'lxml')
    for div in soup.find_all("div", attrs={"class": "address_list"}):
      for li in div.find_all("li", attrs={"class": "left"}):
        ruby = li.find("span").getText()
        url = 'https://www.navitime.co.jp' + li.find("a").get("href")
        name = li.find("a").getText()
        sub_urls.append(url)
        code = url.split("/")[4]
        data = (code, name, ruby, prefecture_code)
        cities.append(data)

    time.sleep(1)

  df = pd.DataFrame(cities, columns=['code', 'name', 'ruby', 'prefecture_code'])
  df.to_csv("data/cities.csv", index=False)
  print(f"cities end : {now()}")

  towns(sub_urls)

def towns(urls: list):
  print(f"towns start : {now()}")
  sub_urls = []
  towns = []
  for url in urls:
    city_code =  url.split("/")[4]
    print(f"----- city_code={city_code} start : {now()} -----")
    # if not city_code == '13113':
    #   continue
    response = requests.get(url)
    soup = BeautifulSoup(response.text, 'lxml')
    for div in soup.find_all("div", attrs={"class": "address_list"}):
      for li in div.find_all("li", attrs={"class": "left"}):
        ruby = li.find("span").getText()
        url = 'https://www.navitime.co.jp' + li.find("a").get("href")
        name = li.find("a").getText()
        sub_urls.append(url)
        code = url.split("/")[4]
        data = (code, name, ruby, city_code)
        towns.append(data)

    time.sleep(1)

  df = pd.DataFrame(towns, columns=['code', 'name', 'ruby', 'city_code'])
  df.to_csv("data/towns.csv", index=False)
  print(f"towns end : {now()}")

def now():
  return datetime.now().strftime("%Y/%m/%d %H:%M:%S")

if __name__ == "__main__":
  crawl()


from flask_restful import fields, marshal_with, reqparse, Resource
import random

def percentage(value):
    '''Return the value if it is a valid percentage. Else raise an exception.'''
    if value >= 0.0 and value <= 50.0:
        return value
    else:
        raise ValueError('{} is not a valid percentage'.format(value))

post_parser = reqparse.RequestParser()

post_parser.add_argument(
    'name', dest='name',
    type=str, required=True,
    help='The unique id of the baverage'
)

post_parser.add_argument(
    'alcohol_percentage', dest='alcohol_percentage',
    type=percentage, default=0.0,
    help='alcohol percentage of the baverage'
)

baverage_fields = {
    'id': fields.Integer,
    'name': fields.String,
    'alcohol_percentage': fields.Float
}

def create_baverage(name, alcohol_percentage):
    return { 'id': random.randrange(1, 1500),
             'name': name,
             'alcohol_percentage': alcohol_percentage
    }

class Baverage(Resource):

    def get(self, id):
        return {'id': id,
                'name': 'coffee',
                'alcohol_percentage': 0.0
        }

    # marshal is pythons' internal serialization
    @marshal_with(baverage_fields)
    def post(self):
        args = post_parser.parse_args() # cleaned
        baverage = create_baverage(args.name, args.alcohol_percentage)
        return baverage


import sys, os, time
from Tkinter import *
from ttk import Frame, Button, Style
import numpy as np
import functions as hfunc
import classes as hclass
import mathematics as hmath
from gui_functions import *
from gui_tetra import *
from gui_umbrella import *

####################################################################################################
## GUI for Initialization
####################################################################################################

class Intro(Frame):

	def __init__(self, parent):
		Frame.__init__(self, parent)
		self.parent = parent
		self.initUI()
		self.select_config()

	def initUI(self):
		self.parent.title("OPTICAL LATTICE")
		self.style= Style()
		self.style.theme_use("default")
		self.pack(fill=BOTH, expand=1)

	def select_config(self):

		def run_standard():
			root = Tk()
			standard = Standard(root)
			root.mainloop()

		def run_umbrella():
			root = Tk()
			umbrella = Umbrella(root)
			root.mainloop()

		ref1 = Label(self)
		ref2 = Label(self)
		ref3 = Label(self)
		ref4 = Label(self)
		owner = Label(self, text="Developer: Hoseong Asher Lee")
		univ = Label(self, text="Miami University, Oxford")
		git = Label(self, text="https://github.com/ohpyupi/optlattices")

		label = Label(self, text=" CHOOSE YOUR LASER CONFIGURATION")
		standard = Button(self, text="Standard Tetrahedron", command=run_standard)
		umbrella = Button(self, text="Umbrella", command=run_umbrella)
		quit = Button(self, text="Quit", command=self.quit)

		label.grid(row=0, column=0)
		ref1.grid(row=1, column=0)
		standard.grid(row=2, column=0)
		ref2.grid(row=3, column=0)
		umbrella.grid(row=4, column=0)

		ref3.grid(row=5, column=0)
		quit.grid(row=6, column=0)
		ref4.grid(row=7)
		owner.grid(row=8)
		univ.grid(row=9)
		git.grid(row=10)


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

import argparse
import json
import os
from os import path
import numpy as np
import scipy as sp
import scipy.sparse as smat
import ctypes
from ctypes import *
from sklearn.preprocessing import normalize as sk_normalize

from xbert.rf_util import PyMatrix, fillprototype, load_dynamic_library

class RandomProject(object):

    """Encode and decode a label into a K-way D-dimensional code.

    feat_mat: L by P matrix
    L: number of label
    P: label feature dimension
    K: range of the code = {0, 1, 2, ..., K-1}
    D: length of the codes for each label (number of hashing functions)
    """

    def __init__(self, feat_mat, kdim, depth, algo, seed):
        self.feat_mat = feat_mat
        self.code_delim = " | "
        self.K, self.D = kdim, depth
        self.L, self.P = feat_mat.shape
        self.algo = algo
        self.random_matrix = np.random.randn(self.P, self.D)

    def get_codes(self):
        if self.algo == 2:  # ordinal
            Z_quant = self.ordinal_quantization(self.feat_mat)
        elif self.algo == 3:  # uniform
            Z_quant = self.uniform_quantization(self.feat_mat)
        else:
            raise NotImplementedError("unknown algo {}".format(self.algo))
        self.hash_func = np.array([self.K ** d for d in range(0, self.D)])
        return Z_quant.dot(self.hash_func)

    # Z: L by D projected label embedding
    # Z_quant: L by D quantized projected label embeddings
    # quantize it by ordinal ranking
    def ordinal_quantization(self, label_embedding):
        Z = label_embedding.dot(self.random_matrix)
        Z_argsort = np.argsort(Z, axis=0)
        bin_size = math.ceil(self.L * 1.0 / self.K)
        Z_quant = []
        for d in range(self.D):
            rank = np.zeros(self.L, dtype=np.int64)
            rank[Z_argsort[:, d]] = np.arange(self.L)
            quantized_rows = (rank // bin_size).tolist()
            Z_quant.append(quantized_rows)
        Z_quant = np.array(Z_quant).T
        return Z_quant

    # quantize it by min/max of each cols into K bins
    def uniform_quantization(self, label_embedding):
        Z = label_embedding.dot(self.random_matrix)
        Z_quant = []
        for d in range(self.D):
            bins = np.linspace(min(Z[:, d]), max(Z[:, d]), self.K)
            quantized_rows = np.digitize(Z[:, d], bins) - 1  # bits = {0, 1, ..., K-1}
            Z_quant.append(quantized_rows)
        Z_quant = np.array(Z_quant).T
        return Z_quant

    def prepare_coding(self, Z_quant):
        # L dimensional array,
        # each entry is the hash code (row idx)
        self.hash_code_arr = Z_quant.dot(self.hash_func)
        rows, cols, vals = [], [], []
        for l in range(self.L):
            rows.append(self.hash_code_arr[l])
            cols.append(l)
            vals.append(1)

        m = self.K ** self.D
        n = self.L
        M = sp.csr_matrix((vals, (rows, cols)), shape=(m, n))
        self.code2label_mat = M

        code2label_set = {}
        for code in np.nonzero(M.indptr[1:] - M.indptr[:-1])[0]:
            code2label_set[code] = set(M.indices[M.indptr[code] : M.indptr[code + 1]])
        self.code2label_set = code2label_set

class corelib(object):
    def __init__(self, dirname, soname, forced_rebuild=False):
        self.clib_float32 = load_dynamic_library(dirname, soname + "_float32", forced_rebuild=forced_rebuild)
        self.clib_float64 = load_dynamic_library(dirname, soname + "_float64", forced_rebuild=forced_rebuild)
        arg_list = [
            POINTER(PyMatrix),
            c_uint32,
            c_uint32,
            c_int32,
            c_uint32,
            c_int32,
            POINTER(c_uint32),
        ]
        fillprototype(self.clib_float32.get_codes, None, arg_list)
        fillprototype(self.clib_float64.get_codes, None, arg_list)

    def get_codes(self, py_feat_mat, depth, algo, seed, codes, verbose=0, max_iter=10, threads=-1):
        clib = self.clib_float32
        if py_feat_mat.dtype == sp.float64:
            clib = self.clib_float64
            if verbose != 0:
                print("perform float64 computation")
        else:
            clib = self.clib_float32
            if verbose != 0:
                print("perform float32 computation")
        clib.get_codes(
            byref(py_feat_mat), depth, algo, seed, max_iter, threads, codes.ctypes.data_as(POINTER(c_uint32)),
        )

forced_rebuild = False
corelib_path = path.join(path.dirname(path.abspath(__file__)), "corelib/")
soname = "rf_linear"
clib = corelib(corelib_path, soname, forced_rebuild)

# SEmatic-aware Code
class SeC(object):
    def __init__(self, kdim, depth, algo, seed, codes):
        assert kdim == 2
        self.kdim = kdim
        self.depth = depth
        self.algo = algo
        self.seed = seed
        self.codes = codes
        self.indptr = sp.cumsum(sp.bincount(codes + 1, minlength=(self.nr_codes + 1)), dtype=sp.uint64)
        self.indices = sp.argsort(codes * sp.float64(self.nr_elements) + sp.arange(self.nr_elements))

    @property
    def nr_elements(self):
        return len(self.codes)

    @property
    def nr_codes(self):
        return 1 << self.depth

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

    def get_code_for_element(self, eid):
        assert 0 <= eid and eid < self.nr_elements
        return self.codes[eid]

    def get_elements_with_code(self, code):
        assert 0 <= code and code < self.nr_codes
        begin, end = self.indptr[code], self.indptr[code + 1]
        return self.indices[begin:end]

    def get_csc_matrix(self):
        return smat.csc_matrix(
            (sp.ones_like(self.indices, dtype=sp.float64), self.indices, self.indptr), shape=(self.nr_elements, self.nr_codes),
        )

    def print(self):
        print("nr_codes: {}".format(self.nr_codes))
        print("nr_elements: {}".format(self.nr_elements))
        print("algo: {}".format(Indexer.algos[self.algo]))
        for nid in range(self.nr_codes):
            labels = " ".join(map(str, self.get_elements_with_code(nid)))
            print("code({nid}): {labels}".format(nid=nid, labels=labels))

class Indexer(object):
    KMEANS = 0
    KDTREE = 1  # KDTREE with Roound-Robin feature splits
    ORDINAL = 2  # random projection with ordinal quantization
    UNIFORM = 3  # random projection with uniform quantization
    BALANCED_ORDINAL = 4  # random projection with balaced ordinal quantization
    SKMEANS = 5  # Spherical KMEANS
    KDTREE_CYCLIC = 11  # KDTREE with cyclic feature splits( 0,...,0, 1,...,1, 2,...,2)
    PURE_RANDOM = 12  # Random assign cluster for each element

    algos = {v: k for k, v in vars().items() if isinstance(v, int)}

    def __init__(self, feat_mat):
        self.py_feat_mat = PyMatrix.init_from(feat_mat)

    @property
    def feat_mat(self):
        return self.py_feat_mat.buf

    @property
    def nr_labels(self):
        return self.feat_mat.shape[0]

    @staticmethod
    def load_indexed_code(code_path, label_feat):
        C = None
        mapping = {
            "none": Indexer.SKMEANS,
            "skmeans": Indexer.SKMEANS,
            "kmeans": Indexer.KMEANS,
            "kdtree": Indexer.KDTREE,
            "random": Indexer.PURE_RANDOM,
            "ordinal": Indexer.BALANCED_ORDINAL,
            "uniform": Indexer.UNIFORM,
        }
        if code_path is None:
            code_path = "none"

        if code_path.lower() in mapping:
            if label_feat is not None:
                algo = mapping[code_path.lower()]
                if algo == Indexer.SKMEANS:
                    label_feat = sk_normalize(label_feat, axis=1, norm="l2", copy=False)
                indexer = Indexer(label_feat)
                code = indexer.gen(
                    kdim=2,
                    depth=indexer.estimate_depth_with_cluster_size(100),
                    algo=algo,
                    seed=0,
                    max_iter=20,
                    threads=1,
                )
                C = code.get_csc_matrix()
        else:
            if code_path.endswith(".npz") and path.exists(code_path):
                C = smat.load_npz(code_path)
            elif path.isdir(code_path) and path.exists(path.join(code_path, "code.npz")):
                C = smat.load_npz(path.join(code_path, "code.npz"))
            else:
                assert False, f"'{code_path}' does not exist. Valid ones {mapping.keys()}"
        return C

    def estimate_depth_with_nr_clusters(self, nr_clusters):
        depth = int(sp.log2(nr_clusters))
        return depth

    def estimate_depth_with_cluster_size(self, cluster_size):
        return self.estimate_depth_with_nr_clusters(self.nr_labels // cluster_size + 1)

    def ordinal_gen(self, kdim, depth, seed):
        sp.random.seed(seed)
        random_matrix = sp.randn(self.feat_mat.shape[1], depth)
        X = self.feat_mat.dot(random_matrix)
        m = self.feat_mat.shape[0] // kdim + [1, 0][self.feat_mat.shape[0] % kdim != 0]
        X = sp.argsort(sp.argsort(X, axis=0), axis=0) // m
        print(X)
        codes = sp.array((X * (kdim ** sp.arange(depth)).reshape(1, -1)).sum(axis=1), dtype=sp.uint32)
        return codes

    def balaced_ordinal_gen(self, kdim, depth, seed, threads=1):
        assert int(2 ** sp.log2(kdim)) == kdim
        sp.random.seed(seed)
        random_matrix = sp.randn(self.feat_mat.shape[1], depth)
        X = PyMatrix(self.feat_mat.dot(random_matrix))
        codes = sp.zeros(X.rows, dtype=sp.uint32)
        new_depth = depth * int(sp.log2(kdim))
        clib.get_codes(X, new_depth, Indexer.KDTREE_CYCLIC, seed, codes, threads=threads)
        return codes

    def gen(self, kdim, depth, algo, seed, max_iter=10, threads=1):
        assert algo in [
            Indexer.KMEANS,
            Indexer.KDTREE,
            Indexer.ORDINAL,
            Indexer.UNIFORM,
            Indexer.BALANCED_ORDINAL,
            Indexer.KDTREE_CYCLIC,
            Indexer.SKMEANS,
            Indexer.PURE_RANDOM,
        ]
        if algo in [
            Indexer.KMEANS,
            Indexer.KDTREE,
            Indexer.KDTREE_CYCLIC,
            Indexer.SKMEANS,
        ]:
            feat_mat = self.py_feat_mat
            codes = sp.zeros(feat_mat.rows, dtype=sp.uint32)
            clib.get_codes(feat_mat, depth, algo, seed, codes, max_iter=max_iter, threads=threads)
        elif algo in [Indexer.ORDINAL, Indexer.UNIFORM]:
            rp_clf = RandomProject(self.feat_mat, kdim, depth, algo, seed)
            codes = rp_clf.get_codes()
        elif algo in [Indexer.BALANCED_ORDINAL]:
            assert int(2 ** sp.log2(kdim)) == kdim
            codes = self.balaced_ordinal_gen(kdim, depth, seed, threads=threads)
        elif algo in [Indexer.PURE_RANDOM]:
            feat_mat = self.py_feat_mat
            codes = sp.zeros(feat_mat.rows, dtype=sp.uint32)
            cluster_size = kdim ** depth
            for idx in range(feat_mat.rows):
                codes[idx] = np.random.randint(0, cluster_size)
        else:
            raise NotImplementedError("unknown algo {}".format(algo))
        return SeC(kdim, depth, algo, seed, codes)

def run_test(data_folder="./datasets/Eurlex-4K"):
    import xbert.rf_linear as rf_linear

    data = rf_linear.Data.load(data_folder, label_emb=None)
    L = smat.load_npz(data_folder + "/L.pifa.npz")
    code = Indexer(L).gen(kdim=2, depth=6, algo=0, seed=5, max_iter=20, threads=1)
    code.print()
    code = Indexer(L).gen(kdim=2, depth=6, algo=5, seed=5, max_iter=20, threads=1)
    code.print()

def load_feature_matrix(src, dtype=sp.float32):
    if src.endswith(".npz"):
        return smat.load_npz(src).tocsr().astype(dtype)
    elif src.endswith(".npy"):
        return smat.csr_matrix(sp.ascontiguousarray(sp.load(src), dtype=dtype))
    else:
        raise ValueError("src must end with .npz or .npy")

def main(args):
    # set hyper-parameters
    input_feat_path = args.input_feat_path
    kdim = args.kdim
    algo = args.algo
    seed = args.seed
    max_iter = args.max_iter
    threads = args.threads
    output_code_dir = args.output_code_dir

    # load label feature matrix (nr_labels * nr_features)
    if path.exists(input_feat_path):
        feat_mat = load_feature_matrix(input_feat_path)
    else:
        raise ValueError("label embedding path does not exist {}".format(input_feat_path))

    if not path.exists(output_code_dir):
        os.makedirs(output_code_dir, exist_ok=True)

    # Indexing algorithm
    # C: nr_labels x nr_codes, stored in csr sparse matrix
    indexer = Indexer(feat_mat)
    if algo == indexer.SKMEANS:
        feat_mat = sk_normalize(feat_mat, axis=1, norm="l2", copy=False)
    depth = indexer.estimate_depth_with_cluster_size(100)+1
    code = indexer.gen(kdim=kdim, depth=depth, algo=algo, seed=seed, max_iter=max_iter, threads=threads)
    C = code.get_csc_matrix()
    print("INDEXER DEPTH", depth, "C", C.shape)

    # save code and args
    output_code_path = path.join(output_code_dir, "code.npz")
    smat.save_npz("{}".format(output_code_path), C, compressed=False)
    output_config_path = path.join(output_code_dir, "config.json")
    with open(output_config_path, "w") as fout:
        fout.write(json.dumps(vars(args), indent=True))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    ## Required parameters
    parser.add_argument(
        "-i",
        "-L",
        "--input-feat-path",
        type=str,
        required=True,
        default="./datasets/Eurlex-4K/L.pifa.npz",
        help="path to the npz file of input label feature matrix (nr_labels * nr_features, CSR)",
    )
    parser.add_argument(
        "-o",
        "-c",
        "--output-code-dir",
        type=str,
        required=True,
        default="./save_models/Eurlex-4K/indexer/code.npz",
        help="path to the output npz file of indexing codes (nr_labels * nr_codes, CSR)",
    )
    # optional
    parser.add_argument("--algo", type=int, default=5, help="0 for KMEANS 5 for SKMEANS (default 5)")
    parser.add_argument("--seed", type=int, default=0, help="random seed (default 0)")
    parser.add_argument("--kdim", type=int, default=2)
    parser.add_argument("--threads", type=int, default=1)
    parser.add_argument("--max-iter", type=int, default=20)

    args = parser.parse_args()
    print(args)
    main(args)


# montecarlo_parallel.py
#
# Omkar H. Ramachandran
# omkar.ramachandran@colorado.edu
#
# Parellelized montecarlo designed to run on CU's Summit supercomputer
#

import numpy as np
import montecarlo_isometric_headers as head
from astropy.io import fits

Nruns = 50
Nrays1 = 6055
Nrays2 = 3000
Nrays3 = 1000
DEGTORAD = np.pi/180

print("**********Montecarlo_Parallel***********")
print("Simulation Paramters")
print("Nruns=",Nruns)
print("N(E_1)=",Nrays1,",N(E_2)=",Nrays2,",N(E_3)=",Nrays3)

hdul_he = fits.open("data/gll_psch_v07.fit")
data_he = hdul_he[1].data
data_he_lon = data_he.field(4)[abs(data_he.field(4))> 0]*DEGTORAD
data_he_lat = data_he.field(3)[abs(data_he.field(4))> 0]*DEGTORAD

hdul_jpd = fits.open("data/gtselected_exp_map_50_60.fits")

Jpd = (hdul_jpd[0].data)[0]

print(Jpd.sum())

Jpd1 = Jpd / Jpd.sum()

R = np.linspace(2,20,10)

data = []

Q = np.zeros([Nruns,len(R)])
Qnorth = np.zeros([Nruns,len(R)])
Qsouth = np.zeros([Nruns,len(R)])
std = np.zeros([Nruns,len(R)])
stdnorth = np.zeros([Nruns,len(R)])
stdsouth = np.zeros([Nruns,len(R)])

for j in range(Nruns):

    theta1_all,phi1_all,theta2_all,phi2_all,theta3_all,phi3_all = head.genRays(Nrays1,Nrays2,Nrays3,Jpd1)
    print("Masking Data around 1FHL sources")
    for i in range(len(data_he_lon)):
        v = np.array([np.cos(data_he_lat[i])*np.cos(data_he_lon[i]),np.cos(data_he_lat[i])*np.sin(data_he_lon[i]),np.sin(data_he_lat[i])])
        # Mask to 2 degrees
        target = np.cos(2*DEGTORAD)

        theta3_new = theta3_all[np.dot(v,np.array([np.cos(theta3_all)*np.sin(phi3_all),np.cos(theta3_all)*np.cos(phi3_all),np.sin(theta3_all)])) - target <= 0]
        phi3_new = phi3_all[np.dot(v,np.array([np.cos(theta3_all)*np.sin(phi3_all),np.cos(theta3_all)*np.cos(phi3_all),np.sin(theta3_all)])) - target <= 0]
        theta2_new = theta2_all[np.dot(v,np.array([np.cos(theta2_all)*np.sin(phi2_all),np.cos(theta2_all)*np.cos(phi2_all),np.sin(theta2_all)])) - target <= 0]
        phi2_new = phi2_all[np.dot(v,np.array([np.cos(theta2_all)*np.sin(phi2_all),np.cos(theta2_all)*np.cos(phi2_all),np.sin(theta2_all)])) - target <= 0]
        theta1_new = theta1_all[np.dot(v,np.array([np.cos(theta1_all)*np.sin(phi1_all),np.cos(theta1_all)*np.cos(phi1_all),np.sin(theta1_all)])) - target <= 0]
        phi1_new = phi1_all[np.dot(v,np.array([np.cos(theta1_all)*np.sin(phi1_all),np.cos(theta1_all)*np.cos(phi1_all),np.sin(theta1_all)])) - target <= 0]

        theta3_all = theta3_new
        phi3_all = phi3_new
        theta2_all = theta2_new
        phi2_all = phi2_new
        theta1_all = theta1_new
        phi1_all = phi1_new

    phi3 = phi3_all[abs(theta3_all) > 80*DEGTORAD]
    theta3 = theta3_all[abs(theta3_all) > 80*DEGTORAD]

    Q[j],std[j] = head.computeQ(R,theta1_all,phi1_all,theta2_all,phi2_all,theta3,phi3)

    phi3 = phi3_all[(theta3_all) > 80*DEGTORAD]
    theta3 = theta3_all[(theta3_all) > 80*DEGTORAD]

    Qnorth[j],stdnorth[j] = head.computeQ(R,theta1_all,phi1_all,theta2_all,phi2_all,theta3,phi3)

    phi3 = phi3_all[-(theta3_all) > 80*DEGTORAD]
    theta3 = theta3_all[-(theta3_all) > 80*DEGTORAD]

    Qsouth[j],stdsouth[j] = head.computeQ(R,theta1_all,phi1_all,theta2_all,phi2_all,theta3,phi3)

Qmean = np.zeros(len(R))
stdmean = np.zeros(len(R))
QmeanN = np.zeros(len(R))
stdmeanN = np.zeros(len(R))
QmeanS = np.zeros(len(R))
stdmeanS = np.zeros(len(R))
for j in range(len(R)):
    Qmean[j] = np.mean(Q[:,j])
    stdmean[j] = np.sqrt(np.mean(std[:,j]**2))
    QmeanN[j] = np.mean(Qnorth[:,j])
    stdmeanN[j] = np.sqrt(np.mean(stdnorth[:,j]**2))
    QmeanS[j] = np.mean(Qsouth[:,j])
    stdmeanS[j] = np.sqrt(np.mean(stdsouth[:,j]**2))


import numpy as np
import h5py as h5
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.cm as cm

prefix = "/media/user/9EAEE48CAEE45DF1/cyl_temp/p-bed-data/1e4/"
prefix = "../../post/"
prefix = "../"
prefix = prefix + "test_pbed_r1_"
ppy =  131
R = 8
num = np.arange(0,730+1)
Dm = 1.15e-6;

name = prefix + str(num[0]).zfill(4) + ".h5"
print(name)
f = h5.File(name)
PX = np.array(f['RWPposition'])
px = PX[0:-2:3]
py = PX[1:-1:3]
Ga = np.array(f['Gamma'])
Nx = int(f['Nx'][0])
Ny = int(f['Ny'][0])
ga = Ga.reshape((Ny,Nx))
ga[np.where(ga>1.0)] = 1.0
gga = np.sum(ga,axis=1)
# print(gga)
print(ppy-1,gga[ppy-1])
print(ppy,gga[ppy])

layer = np.array([ppy-6*(R+2),ppy-4*(R+2),ppy-2*(R+2),ppy])
layer = np.array([66,88,110,131]);
M0 = np.zeros((len(layer))) + np.nan
kkk = []
for i in range(len(layer)):
	kkk.append(np.where(py<=layer[i]))
	M0[i] = np.size(np.where(py<=layer[i]))
colors = cm.rainbow(np.linspace(0,1,len(kkk)))

print('initial num ',M0)

nn = np.zeros(len(num)) + np.nan
M = np.zeros((len(num),len(layer))) + np.nan
for i in range(len(num)):
	nn[i] = i
	name = prefix + str(num[i]).zfill(4) + ".h5"
	print("start process ", name)
	f = h5.File(name)
	PX = np.array(f['RWPposition'])
	px = PX[0:-2:3]
	py = PX[1:-1:3]

	for ii in range(len(layer)):
		M[i,ii] = float(np.size(np.where(py[kkk[ii]]<=layer[ii])))/M0[ii]

k1 = np.zeros(len(layer)) + np.nan
Vs = np.zeros(len(layer)) + np.nan
Vt = np.zeros(len(layer)) + np.nan
Vv = np.zeros(len(layer)) + np.nan
As = np.zeros(len(layer)) + np.nan

tt = nn*1e2
for ii,c in zip(range(len(layer)),colors):
	plt.plot(np.sqrt(tt),M[:,ii],'o',color=c,label=layer[ii])
	Z = np.polyfit(np.sqrt(tt),M[:,ii],1)
	k1[ii] = -Z[0]
	plt.plot(np.sqrt(tt),Z[0]*np.sqrt(tt)+Z[1],color=c)
	Vs[ii] = np.sum(np.sum(ga[1:layer[ii],:]))
	Vt[ii] = (layer[ii]*1.0-1.0)*Nx*1.0
	Vv[ii] = Vt[ii] - Vs[ii]
	As[ii] = Nx - np.sum(ga[layer[ii],:])

plt.legend()
plt.savefig('M.png',dpi=500)
plt.clf()

Deff = np.pi*(0.5*Vv/As*k1)**2/Dm
print('Deff/Dm ', Deff)
plt.plot(Deff,layer,'*')
plt.ylabel('Depth')
plt.xlabel('Deff/Dm')

plt.savefig('Deff',dpi=500)
plt.clf()

f.close()

with open("Deff_1e4.txt","w") as f:
	f.write(str(Dm)+ "\t Dm\n")
	for l,d in zip(layer, Deff):
		f.write(str(l)+" "+str(d)+"\t pos, Deff/Dm\n")


### 003 load data from csv file
import numpy as np
import pandas as pd

print('### 003')
df = pd.read_csv('./movie_data.csv')
print(df.head(3))
print()

### 004 transform words into feature vectors
print('### 004')
from sklearn.feature_extraction.text import CountVectorizer

# 1-gram representation is used by default, we could switch to a 2-gram
# representation by initializing a new CountVectorizer instance with
# ngram_range=(2,2).
count = CountVectorizer()
docs = np.array([
       'The sun is shining',
       'The weather is sweet',
       'The sun is shining and the weather is sweet'])
bag = count.fit_transform(docs)
print(count.vocabulary_)
# feature vectors, are also called the raw term frequencies: tf(t,d)
# the number of times a term t occurs in a document d
print(bag.toarray())
print()

### 005 td-idf
print('### 005')
from sklearn.feature_extraction.text import TfidfTransformer

tfidf = TfidfTransformer()
np.set_printoptions(precision=2)
print(tfidf.fit_transform(count.fit_transform(docs)).toarray())
print()

### 006 clean text data by regexp
print('### 006')
print(df.loc[0, 'review'][-50:])

import re

def preprocessor(text):
    text = re.sub('<[^>]*>', '', text)
    emoticons = re.findall('(?::|;|=)(?:-)?(?:\)|\(|D|P)', text)
    # Although the addition of the emoticon characters to the end of the cleaned document
    # strings may not look like the most elegant approach, the order of the words doesn't
    # matter in our bag-of-words model if our vocabulary only consists of 1-word tokens.
    text = re.sub('[\W]+', ' ', text.lower()) + ' '.join(emoticons).replace('-', '')
    return text

print(preprocessor(df.loc[0, 'review'][-50:]))
print(preprocessor("</a>This :) is :( a test :-)!"))
print()

# now apply our preprocessor function to all movie reviews in our DataFrame
df['review'] = df['review'].apply(preprocessor)

### 007 process documents into tokens
print('### 007')

# One way to tokenize documents is to split them into individual words
# by splitting the cleaned document at its whitespace characters
def tokenizer(text):
    return text.split()

print(tokenizer('runners like running and thus they run'))

# another useful technique is word stemming, which is the process of
# transforming a word into its root form that allows us to map related
# words to the same stem.
from nltk.stem.porter import PorterStemmer

porter = PorterStemmer()
def tokenizer_porter(text):
    return [porter.stem(word) for word in text.split()]

print(tokenizer_porter('runners like running and thus they run'))
print()

### 008 stop-word removal
print('### 008')

from nltk.corpus import stopwords

stop = stopwords.words('english')
print([w for w in tokenizer_porter('a runner likes running and runs a lot')[-10:] if w not in stop])
print()

### 009 train a logistic regression model for document classification
print('### 009')
print('This step will take nearly 2 hours running in a desktop with 16G RAM and E5-2670 CPU')
X_train = df.loc[:25000, 'review'].values
y_train = df.loc[:25000, 'sentiment'].values
X_test = df.loc[25000:, 'review'].values
y_test = df.loc[25000:, 'sentiment'].values

from sklearn.model_selection import GridSearchCV
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LogisticRegression
from sklearn.feature_extraction.text import TfidfVectorizer

# TfidVectorizer is equivalent to CountVectorizer followed by TfidfTransformer.
tfidf = TfidfVectorizer(strip_accents=None, lowercase=False, preprocessor=None)
# The first dictionary uses the TfidfVectorizer with its default settings
# (use_idf=True, smooth_idf=True, and norm='l2') to calculate the tf-idfs;
param_grid = [{'vect__ngram_range': [(1,1)],
               'vect__stop_words': [stop, None],
               'vect__tokenizer': [tokenizer, tokenizer_porter],
               'clf__penalty': ['l1', 'l2'],
               'clf__C': [1.0, 10.0, 100.0]},
              {'vect__ngram_range': [(1,1)],
               'vect__stop_words': [stop, None],
               'vect__tokenizer': [tokenizer, tokenizer_porter],
               'vect__use_idf':[False],
               'vect__norm':[None],
               'clf__penalty': ['l1', 'l2'],
               'clf__C': [1.0, 10.0, 100.0]}
             ]
lr_tfidf = Pipeline([('vect', tfidf), ('clf', LogisticRegression(random_state=0))])
gs_lr_tfidf = GridSearchCV(lr_tfidf, param_grid, scoring='accuracy', cv=5, verbose=1, n_jobs=-1)
gs_lr_tfidf.fit(X_train, y_train)

print('Best parameter set: %s ' % gs_lr_tfidf.best_params_)
print('CV Accuracy: %.3f' % gs_lr_tfidf.best_score_)
clf = gs_lr_tfidf.best_estimator_
print('Test Accuracy: %.3f' % clf.score(X_test, y_test))