# -*- coding: utf-8 -*-
"""
Created on Wed Apr 11 11:26:23 2018

@author: jglasskatz
"""

import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import colors as c
from mpl_toolkits.basemap import Basemap, shiftgrid
import numpy as np
import os
import scipy.stats as st
import matplotlib.animation

#Read in the data
#go to appropriate directory
os.chdir('C:\\Users\\jglasskatz\\Desktop\\CNRA_DATA')

#make csvs df
gwm = pd.read_csv('continuousgroundwatermonthly.csv')
stat = pd.read_csv('gwlstations.csv')
#merge to get relevant data set
gw = gwm.merge(stat, how ='left', on = 'STATION')
#Relevant depths
gw = gw[:][gw['GSE_WSE']>0]
#Get a bunch of different data sets with unique dates
Utime =pd.unique(gw['MSMT_DATE'])
datedata  =[]
for time in Utime:
    df = gw[:][gw['MSMT_DATE']==time]
    datedata.append(df)

#make a list of the ones with 100 samples or more
dfs = []
for df in datedata:
    #print(df.shape)
    if(df.shape[0]>100):
        dfs.append(df)
lons = []
lats = []
depths = []
dates = []
os.chdir('C:\\Users\\jglasskatz\\Desktop\\Depth_Im')
for i in range(0, len(dfs)):
#Do a tester
    test = dfs[i]
    test2 = test.reset_index()
    #They are the same so all unique)

    #Get data
    x1 = test2['LONGITUDE'][:]
    y1 = test2['LATITUDE'][:]
    depth = test2['GSE_WSE'][:]

    fig = plt.figure(figsize=[10,10])
    # specify (nrows, ncols, axnum)
    ax = fig.add_subplot(1, 1, 1)
    #Title
    ax.set_title('Groundwater Depth for ' + test2['MSMT_DATE'][0], fontsize = 14)

    map = Basemap(projection='cyl',llcrnrlat=32.5,urcrnrlat=42.5,llcrnrlon=-125,urcrnrlon=-115,resolution='c', ax=ax)
    #fill it in a bit
    map.drawcoastlines()
    map.drawmapboundary (fill_color='aqua', zorder =1)
    map.fillcontinents (color='#cc9955', zorder =2)
    map.drawstates(zorder =3)
    lon, lat = map(x1,y1)
    lons.append(lon)
    lats.append(lat)
    dates.append(test2['MSMT_DATE'][0])
    depths.append(depth)
    cs = map.scatter(lon, lat, c=depth, cmap='gnuplot', zorder=4)
    fig.colorbar(cs, orientation = 'horizontal')
    #fig.savefig(test2['MSMT_DATE'][0]+'.png')

#################################
#################################
#Make an animation
#The figure to be animated
fig = plt.figure(figsize=[8,8])
# specify (nrows, ncols, axnum)
ax = fig.add_subplot(1, 1, 1)
ax.set_title('Groundwater Depth for ' + dates[0], fontsize = 14)

map = Basemap(projection='cyl',llcrnrlat=32.5,urcrnrlat=42.5,llcrnrlon=-125,urcrnrlon=-115,resolution='c', ax=ax)
    #fill it in a bit
map.drawcoastlines()
map.drawmapboundary (fill_color='aqua', zorder =1)
map.fillcontinents (color='#cc9955', zorder =2)
map.drawstates(zorder =3)
cs = map.scatter(lons[0], lats[0], c=depths[0], cmap='gnuplot', zorder=4)
fig.colorbar(cs, orientation = 'horizontal')

def update(i):
    ax.set_title('Groundwater Depth for ' + dates[i], fontsize = 14)
    cs = map.scatter(lons[i], lats[i], c=depths[i], cmap='gnuplot', zorder=4)

ani = matplotlib.animation.FuncAnimation(fig, update, frames=len(dates), repeat=True)
ani.save('GWEbyMonth.mp4', writer=writer)

#Zoomed in version
#map = Basemap(projection='cyl',llcrnrlat=min(y1)-.2,urcrnrlat=max(y1)+.2,llcrnrlon=min(x1)-1,urcrnrlon=max(x1)+1,resolution='c', ax=ax)


import pywhatkit as kit

def send_message(msg_list, phone_number):
    if msg_list:
        concantenate_msg = ""
        for msg in msg_list:
            concantenate_msg += f"Date: {msg['date']} | Bank: {msg['url']} | Country: {msg['country']} | Value: {msg['value']}\n\n"
        try:
            kit.sendwhatmsg_instantly(
                phone_no=phone_number,
                message=concantenate_msg,
                wait_time=15,
                tab_close=True,
                close_time=1
            )
            print("All messages sent successfully!")
        except Exception as e:
            print(e)
            return
    else:
        print("No messages to send.")


import numpy as np

from .replay_buffer import ReplayBuffer

# Reference: https://github.com/LeejwUniverse/following_deepmid/tree/master/jungwoolee_pytorch/100%20Algorithm_For_RL/01%20sum_tree
class PERBuffer(ReplayBuffer):
    def __init__(self, buffer_size, uniform_sample_prob=1e-3):
        super(PERBuffer, self).__init__(buffer_size)
        self.tree_size = (self.buffer_size * 2) - 1
        self.first_leaf_index = self.buffer_size - 1

        self.sum_tree = np.zeros(self.tree_size)  # define sum tree
        self.tree_index = self.first_leaf_index  # define sum_tree leaf node index.

        self.max_priority = 1.0
        self.uniform_sample_prob = uniform_sample_prob

    def store(self, transitions):
        if self.first_store:
            self.check_dim(transitions[0])

        for transition in transitions:
            self.buffer[self.buffer_index] = transition
            new_priority = (
                transition["priority"]
                if "priority" in transition
                else self.max_priority
            )
            self.add_tree_data(new_priority)

            self.buffer_counter = min(self.buffer_counter + 1, self.buffer_size)
            self.buffer_index = (self.buffer_index + 1) % self.buffer_size

    def add_tree_data(self, new_priority):
        self.update_priority(new_priority, self.tree_index)

        self.tree_index += 1  # count current sum_tree index
        if self.tree_index == self.tree_size:  # if sum tree index achive last index.
            self.tree_index = self.first_leaf_index  # change frist leaf node index.

    def update_priority(self, new_priority, index):
        ex_priority = self.sum_tree[index]
        delta_priority = new_priority - ex_priority
        self.sum_tree[index] = new_priority
        self.update_tree(index, delta_priority)

        self.max_priority = max(self.max_priority, new_priority)

    def update_tree(self, index, delta_priority):
        # index is a starting leaf node point.
        while index > 0:
            index = (index - 1) // 2  # parent node index.
            self.sum_tree[index] += delta_priority

    def search_tree(self, num):
        index = 0  # always start from root index.
        while index < self.first_leaf_index:
            left = (index * 2) + 1
            right = (index * 2) + 2

            if num <= self.sum_tree[left]:  # if child left node is over current value.
                index = left  # go to the left direction.
            else:
                num -= self.sum_tree[left]  # if child left node is under current value.
                index = right  # go to the right direction.

        return index

    def sample(self, beta, batch_size):
        assert self.sum_tree[0] > 0.0
        uniform_sampling = np.random.uniform(size=batch_size) < self.uniform_sample_prob
        uniform_size = np.sum(uniform_sampling)
        prioritized_size = batch_size - uniform_size

        uniform_indices = list(
            np.random.randint(self.buffer_counter, size=uniform_size)
            + self.first_leaf_index
        )

        targets = np.random.uniform(size=prioritized_size) * self.sum_tree[0]
        prioritized_indices = [self.search_tree(target) for target in targets]

        indices = np.asarray(uniform_indices + prioritized_indices)
        priorities = np.asarray([self.sum_tree[index] for index in indices])
        assert len(indices) == len(priorities) == batch_size

        uniform_probs = np.asarray(1.0 / self.buffer_counter)
        prioritized_probs = priorities / self.sum_tree[0]

        usp = self.uniform_sample_prob
        sample_probs = (1.0 - usp) * prioritized_probs + usp * uniform_probs
        weights = (uniform_probs / sample_probs) ** beta
        weights /= np.max(weights)
        batch = [self.buffer[idx] for idx in indices - self.first_leaf_index]

        transitions = self.stack_transition(batch)

        sampled_p = np.mean(priorities)
        mean_p = self.sum_tree[0] / self.buffer_counter
        return transitions, weights, indices, sampled_p, mean_p

    @property
    def size(self):
        return self.buffer_counter


from tkinter import *
from tkinter import filedialog
from tkinter.messagebox import showwarning, showinfo
from tkinter.ttk import *
from openpyxl import Workbook
from openpyxl.utils import get_column_letter
import csv

class GUI():
    def __init__(self):
        self.window = Tk()

        self.window.title("Μετατροπή ενός αρχείου csv σε xlsx")
        self.window.resizable(False, False)
        self.create_widgets()

    def getInputFilename(self):
        fName = filedialog.askopenfilename(initialdir=".", title="Επιλέξτε το αρχείο csv",
                                           filetypes=(("csv files", "*.csv"), ("all files", "*.*")))

        if fName == "":
            return

        self.btnOpenInputFile.configure(state='disabled')
        self.inputFilename.set(fName)

        self.data = self.parseCsvData(fName)
        self.btnRun.configure(state='normal')

    def parseCsvData(self, inputFile):
        data = list()

        try:
            with open(inputFile, 'rt', encoding='utf-8-sig') as f:
                dialect = csv.Sniffer().sniff(f.readline())
                f.seek(0)

                reader = csv.reader(f, delimiter=dialect.delimiter, quotechar=dialect.quotechar)

                for row in reader:
                    data.append(row)
        except:
            with open(inputFile, 'rt') as f:
                dialect = csv.Sniffer().sniff(f.readline())
                f.seek(0)

                reader = csv.reader(f, delimiter=dialect.delimiter, quotechar=dialect.quotechar)

                for row in reader:
                    data.append(row)

        return data

    def saveFile(self, data, outputFile):
        wb = Workbook()
        ws = wb.active

        for entry in data:
            ws.append(entry)

        column_widths = []
        for row in ws.iter_rows():
            for i, cell in enumerate(row):
                try:
                    column_widths[i] = max(column_widths[i], len(str(cell.value)))
                except IndexError:
                    column_widths.append(len(str(cell.value)))

        for i, column_width in enumerate(column_widths):
            ws.column_dimensions[get_column_letter(i + 1)].width = column_width * 1.23

        notSaved = True

        while notSaved:
            try:
                wb.save(outputFile)
            except:
                showwarning(title="Αρχείο σε χρήση...",
                            message="Παρακαλώ κλείστε το αρχείο '{}' ώστε να ολοκληρωθεί η αποθήκευση.".format(
                                outputFile))
            else:
                notSaved = False

    def run(self):
        self.btnRun.configure(state='disabled')

        inputFile = self.inputFilename.get()
        outputFile = inputFile.replace(".csv", ".xlsx")

        self.saveFile(self.data, outputFile)

        showinfo(title='Ολοκλήρωση Εκτέλεσης',
                 message=f'Η μετατροπή του αρχείου csv σε xlsx ολοκληρώθηκε.')

    def create_widgets(self):
        self.fData = Frame(self.window)

        self.lInputFile = Label(self.fData, text="Αρχείο csv:")
        self.lInputFile.grid(column=0, row=0, padx=10, pady=10, sticky=E)

        self.inputFilename = StringVar()
        self.inputFilename.set('')
        self.ntrInputFilename = Entry(self.fData, width=128, state='readonly', textvariable=self.inputFilename)
        self.ntrInputFilename.grid(column=1, row=0, padx=10, pady=10, sticky=W)

        self.btnOpenInputFile = Button(self.fData, text="Επιλέξτε αρχείο...", command=self.getInputFilename)
        self.btnOpenInputFile.grid(column=2, row=0, padx=10, pady=10)

        self.btnRun = Button(self.fData, text="Εκτέλεση", command=self.run, state='disabled')
        self.btnRun.grid(column=1, row=10, padx=10, pady=10)

        self.fData.pack()

gui = GUI()
gui.window.mainloop()


import os
from glob import glob
import shutil
import pprint

import traceback
import warnings
import json
import yaml

import intake

import dask
import xarray as xr

from toolz import curry

from . config import cache_catalog_dir, cache_catalog_prefix
from . registry import _DERIVED_VAR_REGISTRY, _QUERY_DEPENDENT_OP_REGISTRY

os.makedirs(cache_catalog_dir, exist_ok=True)

class Collection(object):
    """
    Catalog + Query + Operation = Unique Dataset

    Extend an intake-esm catalog with the ability to:
      (1) Compute derived variables combining multipe catalog entries;
      (2) Apply postprocessing operations to returned datasets;
      (3) Optionally cache the results and curate a discrete catalog of processed assets;
          skip computation on subsequent calls for existing cataloged assets.

    Parameters
    ----------

    name : str
      Name of this collection.

    catalog : str
      JSON file defining `intake-esm` catalog

    query : dict
      Dictionary of keyword, value pairs defining a query subsetting the `catalog`.

    postproccess : callable, list of callables, optional
      Call these functions on each dataset following aggregation.

    postproccess_kwargs: dict, list of dict's, optional
       Keyword arguments to `postprocess` functions.
       **Note** These arguments must be serializable to yaml.

    persist : bool, optional
       Persist datasets to disk.

    cache_dir : str, optional
       Directory for storing cache files.

    cache_format : str
       Format in which to store cache files; can be "nc" or "zarr"

    kwargs : dict
       Defaults for keyword arguments to pass to `intake_esm.core.esm_datastore.to_dataset_dict`.
       (Note: these can be updated/overridden for each call to self.dsets below.)
    """
    def __init__(self,
                 name,
                 esm_collection_json,
                 query,
                 postproccess=[],
                 postproccess_kwargs=[],
                 persist=False,
                 cache_dir='.',
                 cache_format='nc',
                 **kwargs,
                ):

        self.name = name

        # TODO: what should we do if "variable" is in the query?
        #       for now, just remove it
        variable = query.get('variable', None)
        if variable is not None:
            warnings.warn('removed "variable" key from query')

        # TODO: accept multiple catalogs and concatenate them into one?
        self.catalog = intake.open_esm_datastore(esm_collection_json).search(**query)

        # setup cache info
        assert cache_format in ['nc', 'zarr'], f'unsupported format {cache_format}'
        self._format = cache_format
        self.persist = persist
        self.cache_dir = cache_dir
        if self.persist and not os.path.exists(cache_dir):
            os.makedirs(cache_dir)
        self._open_cache_kwargs = dict() # TODO: get this from config?
        if self._format == 'zarr' and 'consolidated' not in self._open_cache_kwargs:
            self._open_cache_kwargs['consolidated'] = True

        self._to_dsets_kwargs_defaults = kwargs

        # setup operators attrs
        if not isinstance(postproccess, list):
            postproccess = [postproccess]
        self.operators = postproccess

        if not postproccess_kwargs:
            self.ops_kwargs = [{} for op in self.operators]
        else:
            if not isinstance(postproccess_kwargs, list):
                postproccess_kwargs = [postproccess_kwargs]
            assert len(postproccess_kwargs) == len(postproccess), 'mismatched ops/ops_kwargs'
            self.ops_kwargs = postproccess_kwargs

        # pull out "prepocess"
        preprocess_name = None if 'preprocess' not in kwargs else kwargs['preprocess'].__name__

        # build origins dict
        # TODO: save more about the function, including the source code
        #       rather than simply the name
        self.origins_dict = dict(
            name=name,
            esm_collection=esm_collection_json,
            query=query,
            preprocess=preprocess_name,
            operators=[op.__name__ for op in self.operators],
            operator_kwargs=self.ops_kwargs,
        )

        # assemble database of existing cache's
        self._assemble_cache_db()

    def _assemble_cache_db(self):
        """loop over yaml files in cache dir; find ones that match"""
        self._cache_files = {}
        if not self.persist:
            return

        cache_id_files = self._find_cache_id_files()

        if cache_id_files:
            for file in cache_id_files:
                try:
                    with open(file, 'r') as fid:
                        cache_id = yaml.load(fid, Loader=yaml.Loader)
                except: # TODO: be more informative here
                    print(f'cannot read cache: {file}...skipping')
                    print('skipping')

                cache_origins_dict = {
                    k: cache_id[k]
                    for k in self.origins_dict.keys()
                    if k in cache_id
                }
                if cache_origins_dict == self.origins_dict:
                    variable = cache_id['variable']
                    key = cache_id['key']
                    if variable not in self._cache_files:
                        self._cache_files[variable] = {}
                    self._cache_files[variable][key] = cache_id['asset']

    def to_dataset_dict(self, variable, compute=True, clobber=False,
                        prefer_derived=False, **kwargs):
        """
        Get dataset_dicts for assets.

        Parameters
        ----------

        variable : str or list
          The variable or list of variables to read and process.

        compute : bool, optional (default=True)
          Specify whether to call compute on the final datasets.

        clobber : bool, optional (default=False)
          Remove and recreate any existing cache.

        kwargs : dict, optional
          Keyword arguments to pass to `intake_esm.core.esm_datastore.to_dataset_dict`

        """

        # TODO: `variable` could be more flexible,
        # i.e. it could be a query dict with a require "variable" key

        if isinstance(variable, list):
            dsets_list = [
                self.to_dataset_dict(v, compute, clobber, prefer_derived, **kwargs)
                for v in variable
            ]
            keys = list(set([k for dsets in dsets_list for k in dsets.keys()]))
            dsets_merged = {}
            for key in keys:
                ds_list = []
                for dsets in dsets_list:
                    if key in dsets:
                        ds_list.append(dsets[key])
                dsets_merged[key] = xr.merge(ds_list)
            return dsets_merged

        # TODO: check for lock file, wait if present
        if self._cache_exists(variable, clobber):
            return self._read_cache(variable)

        else:
            # TODO: set lock file
            to_dsets_kwargs = kwargs.copy()
            to_dsets_kwargs.update(self._to_dsets_kwargs_defaults)
            # TODO: optionally spin up a cluster here
            return self._generate_dsets(
                variable, compute, prefer_derived, **to_dsets_kwargs
            )
            # TODO: release lock

    def _generate_dsets(self, variable, compute, prefer_derived, **kwargs):
        """Do the computation necessary to make `dsets`"""

        # check for variable in catalog
        # TODO: we're handling empty results below:
        #       suppress intake-esm warning or change logic
        catalog_subset_var = self.catalog.search(variable=variable)

        # check for variable in derived registry
        is_derived = variable in _DERIVED_VAR_REGISTRY

        if len(catalog_subset_var):
            if is_derived and not prefer_derived:
                warnings.warn(
                    f'found variable "{variable}" in catalog and derived_var registry'
                )
                is_derived = False
        else:
            if not is_derived:
                raise ValueError(f'variable not found {variable}')

        if is_derived:
            derived_var_obj = _DERIVED_VAR_REGISTRY[variable]
            query_vars = derived_var_obj.dependent_vars
            catalog_subset_var = self.catalog.search(variable=query_vars)

        dsets = catalog_subset_var.to_dataset_dict(**kwargs)
        key_info = _intake_esm_get_keys_info(catalog_subset_var)

        if is_derived:
            for key, ds in dsets.items():
                dsets[key] = derived_var_obj(ds)

        for key in dsets.keys():
            for op, kw in zip(self.operators, self.ops_kwargs):
                if hash(op) in _QUERY_DEPENDENT_OP_REGISTRY:
                    op_obj = _QUERY_DEPENDENT_OP_REGISTRY[hash(op)]
                    dsets[key] = op_obj(
                        dsets[key],
                        query_dict=key_info[key],
                        **kw,
                    )
                else:
                    dsets[key] = op(dsets[key], **kw)

        if compute:
            dsets = {k: ds.compute() for k, ds in dsets.items()}

        if self.persist:
            self._make_cache(dsets, variable)

        return dsets

    def _make_cache(self, dsets, variable):
        """write cache file"""

        cache_files = {variable: {}}
        for key, ds in dsets.items():

            cache_id_dict = self.origins_dict.copy()
            cache_id_dict['key'] = key
            cache_id_dict['variable'] = variable
            cache_id_dict['asset'] = self._gen_cache_file_name(key, variable)

            cache_id_file = self._gen_cache_id_file_name(key, variable)

            if self._format == 'nc':
                ds.to_netcdf(cache_id_dict['asset'])
            elif self._format == 'zarr':
                ds.to_zarr(cache_id_dict['asset'], mode='w', consolidated=True)

            with open(cache_id_file, 'w') as fid:
                yaml.dump(cache_id_dict, fid)

            cache_files[variable][key] = cache_id_dict['asset']

        self._cache_files.update(cache_files)

    def _read_cache(self, variable):
        """read cache files"""
        dsets = {}
        for key, asset in self._cache_files[variable].items():
            if self._format == 'nc':
                with xr.open_dataset(asset, **self._open_cache_kwargs) as ds:
                    dsets[key] = ds
            elif self._format == 'zarr':
                with xr.open_zarr(asset, **self._open_cache_kwargs) as ds:
                    dsets[key] = ds
        return dsets

    def _cache_exists(self, variable, clobber):
        """determine if cache files exist (or clobber them)"""

        if variable in self._cache_files:
            if clobber:
                for asset in self._cache_files[variable].values():
                    if os.path.exists(asset):
                        self._remove_asset(asset)
                return False
            else:
                return all(
                    [os.path.exists(asset) for asset in self._cache_files[variable].values()]
                )
        else:
            return False

    def _remove_asset(self, asset):
        """delete asset from disk"""
        if not os.path.exists(asset):
            return
        if self._format == 'nc':
            os.remove(asset)
        elif self._format == 'zarr':
            shutil.rmtree(asset)

    def _gen_cache_file_name(self, key, variable):
        """generate a file cache name"""
        # TODO: accept a user-provided callable to generate human-readable
        #       file name
        token = dask.base.tokenize(self.origins_dict, key, variable)
        return f'{self.cache_dir}/{token}.{self._format}'

    def _gen_cache_id_file_name(self, key, variable):
        """generate a unique cache file name"""
        token = dask.base.tokenize(self.origins_dict, key, variable)
        return f'{cache_catalog_dir}/{cache_catalog_prefix}-{token}.yml'

    def _find_cache_id_files(self):
        return sorted(glob(f'{cache_catalog_dir}/{cache_catalog_prefix}-*.yml'))

    def __repr__(self):
        return pprint.pformat(self.origins_dict, indent=2, width=1)

class derived_var(object):
    """
    Support computation of variables that depend on multiple variables,
    i.e., "derived vars"
    """
    def __init__(self, dependent_vars, func):
        self.dependent_vars = dependent_vars
        self._callable = func

    def __call__(self, ds, **kwargs):
        """call the function to compute derived var"""
        self._ensure_variables(ds)
        return self._callable(ds, **kwargs)

    def _ensure_variables(self, ds):
        """ensure that required variables are present"""
        missing_var = set(self.dependent_vars) - set(ds.variables)
        if missing_var:
            raise ValueError(f'Variables missing: {missing_var}')

class query_dependent_op(object):
    """
    Support calling functions that depend on the values of the query.
    """
    def __init__(self, func, query_keys):
        self._callable = func
        self._query_keys = query_keys

    def __call__(self, ds, query_dict, **kwargs):
        """call function with query keys added to keyword args"""
        kwargs.update({k: query_dict[k] for k in self._query_keys})
        return self._callable(ds, **kwargs)

@curry
def register_derived_var(func, varname, dependent_vars):
    """register a function for computing derived variables"""
    if varname in _DERIVED_VAR_REGISTRY:
        warnings.warn(
            f'overwriting derived variable "{varname}" definition'
        )

    _DERIVED_VAR_REGISTRY[varname] = derived_var(
        dependent_vars, func,
    )
    return func

@curry
def register_query_dependent_op(func, query_keys):
    """register a function for computing derived variables"""
    func_hash = hash(func)
    if func_hash in _QUERY_DEPENDENT_OP_REGISTRY:
        warnings.warn(
            f'overwriting query dependent operator "{func.__name__}" definition'
        )

    _QUERY_DEPENDENT_OP_REGISTRY[func_hash] = query_dependent_op(
        func, query_keys,
    )
    return func

def _intake_esm_get_keys_info(cat):
    """return a dictionary with the values of components of the keys in an
       intake catalog
    """
    groupby_attrs_values = {
        k: cat.unique(columns=k)[k]['values']
        for k in cat.groupby_attrs
    }
    key_info = {k: {} for k in cat.keys()}
    for key in cat.keys():
        for attr in cat.groupby_attrs:
            values = groupby_attrs_values[attr]
            match = [v for v in values if v in key]
            assert len(match) == 1, 'expecting a unique match'
            key_info[key][attr] = match[0]
    return key_info


import numpy as np
import pandas as pd

import matplotlib.pyplot as plt

if __name__ == "__main__":
    k = 30
    data_df = pd.read_csv(rf"C:\Users\liori\Downloads\wandb_export_stoch_{k}.csv", sep=',')
    possible_n_shots = sorted(data_df["n_shots"].unique())
    possible_n_epochs = sorted(data_df["n_test_epochs"].unique())
    possible_adapt_steps = sorted(data_df["test_adapt_steps"].unique())
    means_df = data_df.groupby(["n_shots", "n_test_epochs", "test_adapt_steps"]).mean()["test_accuracy"].reset_index()
    stderr_df = data_df.groupby(["n_shots", "n_test_epochs", "test_adapt_steps"]).sem()["test_accuracy"].reset_index()
    stderr_df.rename(columns={'test_accuracy': 'test_accuracy_stderr'}, inplace=True)

    full_df = means_df.join(stderr_df.set_index(["n_shots", "n_test_epochs", "test_adapt_steps"]),
                                           on=["n_shots", "n_test_epochs", "test_adapt_steps"])
    full_df.to_csv(rf"C:\Users\liori\Downloads\accs_stoch_{k}.csv", index=False)

    # full_df = pd.read_csv(rf"C:\Users\liori\Downloads\accs_stoch_{k}.csv", sep=',')
    # possible_n_shots = sorted(full_df["n_shots"].unique())
    # possible_n_epochs = sorted(full_df["n_test_epochs"].unique())
    # possible_adapt_steps = sorted(full_df["test_adapt_steps"].unique())

    plt.figure()
    for n_steps in possible_n_epochs:
        df_filtered = full_df[full_df["n_test_epochs"]==n_steps]
        if n_steps == 0:
            for n_epochs in possible_adapt_steps:
                # if n_epochs == 50:
                #      continue
                df_filtered2 = df_filtered[df_filtered["test_adapt_steps"]==n_epochs]
                plt.errorbar(df_filtered2["n_shots"], df_filtered2["test_accuracy"],
                             yerr=df_filtered2["test_accuracy_stderr"], label=f"Meta-testing, {n_epochs} epochs")
        else:
            if n_steps not in [5,40]:
                continue
            plt.errorbar(df_filtered["n_shots"], df_filtered["test_accuracy"],
                         yerr=df_filtered["test_accuracy_stderr"], label=f"{n_steps} adapt steps")
    plt.title(f"Test accuracies vs n_shots")
    #plt.legend(loc="lower right")
    lgd = plt.legend(loc='upper right', bbox_to_anchor=(1.5, 1.05))
    plt.savefig(f"test_accuracies_stoch{k}.png", bbox_extra_artists=(lgd,), bbox_inches='tight')
    plt.close()


import sqlite3
import os
import matplotlib.pyplot as plt
import numpy as np
import json
import requests

def createRatingTable():

    path = os.path.dirname(os.path.abspath(__file__))
    conn = sqlite3.connect(path+'/'+"Movies.db")
    cur = conn.cursor()

    cur.execute("CREATE TABLE IF NOT EXISTS IMDB_Ratings (movie_id INTEGER PRIMARY KEY, MetacriticScore INTEGER, ImdbRate INTEGER, imdb_id TEXT)")

    conn.commit()

    return cur,conn

def retrieveIMDBdata():
    response = requests.get('https://imdb-api.com/API/AdvancedSearch/k_u7vte5gf?title_type=feature&release_date=2019-01-01,2022-01-01&groups=oscar_nominees&count=250&sort=alpha,asc')
    data = json.loads(response.text)
    movie_ratings = []
    for movies in data['results']:
        meta = int(movies['metacriticRating'])
        imdb = int(float(movies['imDbRating'])*10)
        title = movies['title']
        id = movies['id']
        tuple = meta,imdb,title,id
        movie_ratings.append(tuple)
    return movie_ratings

def fillRatingTable(movieratings, cur, conn):
    cur.execute("SELECT max(movie_id) FROM IMDB_Ratings")
    count = cur.fetchone()[0]
    if count ==  None:
        count = 0

    for i in range(count, count+26):
        meta = movieratings[i][0]
        imdb = movieratings[i][1]
        title = movieratings[i][2]
        id = movieratings[i][3]
        cur.execute(f"SELECT movie_id from Movies WHERE movie_title = ?", (title,))
        movieid = cur.fetchone()[0]
        cur.execute("INSERT OR IGNORE INTO IMDB_Ratings (movie_id, MetacriticScore, ImdbRate, imdb_id) VALUES (?,?,?,?)", (movieid, meta, imdb, id))

    conn.commit()

def graph_ratings_movies1to30(cur, conn):
    cur.execute("SELECT Movies.movie_title, IMDB_Ratings.ImdbRate, IMDB_Ratings.MetacriticScore FROM IMDB_Ratings JOIN Movies ON Movies.movie_id = IMDB_Ratings.movie_id")
    res = cur.fetchall()[:30]
    ngroups = len(res)

    imdb = []
    meta = []
    titles = []

    for tups in res:
        title = tups[0]
        titles.append(title)
        imdbrate = tups[1]
        imdb.append(imdbrate)
        metarate = tups[2]
        meta.append(metarate)

    fig, ax = plt.subplots()
    index = np.arange(ngroups)
    bar_width = 0.25
    opacity = 0.8

    rects1 = plt.bar(index, imdb, bar_width, alpha = opacity, color = 'navy', label = "Audience Ratings")
    rects2 = plt.bar(index + bar_width, meta, bar_width, alpha = opacity, color = 'salmon', label = "Critic Ratings")

    plt.xlabel('Movie')
    plt.ylabel('Ratings')
    plt.title('Ratings by Movie')
    plt.xticks(index + bar_width, titles, fontsize = 'small')
    fig.autofmt_xdate()
    plt.legend()

    plt.tight_layout()
    plt.show()

def pieChartofRatings(cur, conn):
    cur.execute("SELECT ImdbRate, MetacriticScore FROM IMDB_Ratings")
    res = cur.fetchall()

    audience_higher = 0
    critic_higher = 0
    tie = 0

    for tups in res:
        audience = tups[0]
        critic = tups[1]
        if audience > critic:
            audience_higher += 1
        if critic > audience:
            critic_higher += 1
        if audience == critic:
            tie += 1

    labels = 'Audience', 'Critic', 'Tie'
    sizes = [audience_higher, critic_higher, tie]
    explode = (0, 0.1, 0)
    colors = ('lightgreen', 'deepskyblue', 'tomato')

    fig1, ax1 = plt.subplots()
    ax1.pie(sizes, explode=explode, labels=labels, autopct='%1.1f%%',shadow=True, startangle=90, colors = colors)
    plt.title("Do Critics or Audience Rate Movies Higher?")

    ax1.axis('equal')

    plt.show()

    sizetup = (audience_higher, critic_higher, tie)
    return(sizetup)

def findAverageRatingDifference(cur, conn):
    cur.execute("SELECT ImdbRate, MetacriticScore FROM IMDB_Ratings")
    res = cur.fetchall()

    total = len(res)

    difference = 0

    for tups in res:
        audience = tups[0]
        critic = tups[1]
        dif = abs(audience - critic)
        difference += dif

    average = difference/total
    return average

def averageRatetoFile(cur, average, audience, critic, tie):

    file = open('ImdbOutfile.txt', 'w')

    cur.execute("SELECT max(movie_id) FROM IMDB_Ratings")
    totalintable = cur.fetchone()[0]

    file.write("Number of times (out of {}) the Audience score was greater than the Critic score: {}\n".format(totalintable, audience))
    file.write("Number of times (out of {}) the Critic score was greater than the Audience score: {}\n".format(totalintable, critic))
    file.write("Number of times (out of {}) the Audience score was the same as the Critic score: {}\n".format(totalintable, tie))

    file.write("Average difference between critic and audience rating scores: {}".format(average))

    file.close()

def main():
    cur,conn = createRatingTable()
    movieratings = retrieveIMDBdata()
    fillRatingTable(movieratings, cur, conn)
    # graph_ratings_movies1to30(cur, conn)
    # audience, critic, tie = pieChartofRatings(cur, conn)
    # average = findAverageRatingDifference(cur, conn)
    # averageRatetoFile(cur, average, audience, critic, tie)

if __name__ == "__main__":
    main()


import tushare as ts
from Entity.Price import Price
from LongtermStrategy import  LongtermStrategy
from ShortTermStrategy import ShortTermStrategy
import logging
import datetime

def main():
    logging.basicConfig(level=logging.DEBUG)

    start = '2016-01-01'
    now = datetime.datetime(2018, 4, 1)
    totalBenefit = 0
    for i in range(1):
        now = now + datetime.timedelta(days=1)
        result = simulate(start, now)
        benefit = sum(result)
        totalBenefit += benefit
        if benefit < 0:
            logging.error(result)
    logging.error('average benfit: %s', totalBenefit/365)

def simulate(start, now):
    p = generateHistoricalData(start, now.strftime('%Y-%m-%d'))
    # commandCenter = LongtermStrategy(p)
    commandCenter = ShortTermStrategy(p)

    while now < datetime.datetime(2019, 4, 1):
        next = now + datetime.timedelta(days=1)
        data = ts.get_hist_data('cyb', start=now.strftime('%Y-%m-%d'), end=next.strftime('%Y-%m-%d'), ktype="5")[::-1]
        for row in data.iterrows():
            open = float(row[1]['open'])
            close = float(row[1]['close'])
            high = float(row[1]['high'])
            low = float(row[1]['low'])
            date = row[0].split(' ')[0]

            currentPrice = Price(high, low, open, close, date)

            decision = commandCenter.do(currentPrice)
            if decision.Type == 'Quit with benefit' or decision.Type == 'First Position' \
                    or decision.Type == "Stop with loss" or decision.Type == 'Follow Buying':
                logging.info(decision)
                logging.info(row)
        now = next
    return commandCenter.account

# Exclusive end Date
def generateHistoricalData(startDate, endDate):
    data = ts.get_hist_data('cyb', start=startDate, end=endDate, ktype="D")
    dailyPrices = []

    for row in data.iterrows():
        open = float(row[1]['open'])
        close = float(row[1]['close'])
        high = float(row[1]['high'])
        low = float(row[1]['low'])
        date = row[0]
        dailyPrices.insert(0, Price(high, low, open, close, date))

    return dailyPrices

if __name__ == '__main__':
    main()


import ccxt
import time

def start_check_xrp():
    """Не используются функции, try except и прочие вещи для ускорения кода, чтобы лишний раз
       не замедляться, не создавать новые объекты, вызовы и пр.(хотя тут для читебальности
       и целостности они бы не помешали"""
    binance = ccxt.binance()
    symbol = 'XRP/USDT'  # пару которую ищем
    check_interval = 60

    while True:
        ticker = binance.fetch_ticker(symbol)  # получаем словарь нашей пары
        price = ticker['last']  # последняя цена
        max_price = ticker['high']  # максимальная

        # Проверка, упала ли цена на 1% от максимальной цены за последний час
        if price < max_price * 0.99:
            print(f"Цена {symbol} упала на 1% от максимальной цены за последний час: {price}")

        time.sleep(check_interval)

if __name__ == "__main__":
    start_check_xrp()


import time

def euler_06 (n : int = 100 ) -> int:
    a = sum(map(lambda x : x**2, range(1, 101)))
    return sum(range(1, 101))**2 - a

start = time.time()
answer = euler_06()
end = time.time()
print(f'Answer: {answer}')
print(f'Time to complete: {end - start:.4f}s\n')


import cv2
import mediapipe as mp
import time
import math

dist0 = 0
def openClose(frame, handLms, w, h):
    global dist0
    distcurr=utility(handLms, w, h)
    if(frame == 1):
        dist0 = distcurr
    if((frame%4 +1) == 1 and abs(handLms.landmark[mp.solutions.hands.HandLandmark.WRIST].x*w - w/2)<w/2 and abs(handLms.landmark[mp.solutions.hands.HandLandmark.WRIST].y*h - h/2)<h/2):
        change = dist0- distcurr
        dist = utility(handLms, w, h)
        # print(change)
        if(change>+50):
            print("Hand is closing ")
            return 1
        elif(change<-50):
            print("Hand is opening ")
            return 2
    return 0

time_left = time.time()
time_right = time.time()
left_taken = False
right_taken = False
def sliderHori(handLms, frame, w=1):
    global time_left
    global time_right
    global left_taken
    global right_taken
    for id, lm in enumerate(handLms.landmark):
        if(id==0):
            cx = int((lm.x)*w)
            if(cx < 80 and left_taken== False):
                time_left = time.time()
                left_taken = True

            if(cx > 200 and right_taken == False):
                time_right = time.time()
                right_taken = True
    if(time_left<time_right and right_taken and left_taken):
        print("Hand is moving LEFT")
        return 3
    if(time_right<time_left and right_taken and left_taken):
        print("Hand is moving RIGHT")
        return 4
    return 0

time_up = time.time()
time_down = time.time()
up_taken = False
down_taken = False
def sliderVer(handLms, frame, h=1):
    global time_up
    global time_down
    global up_taken
    global down_taken
    for id, lm in enumerate(handLms.landmark):
        if(id==0):
            cy = int((lm.y)*h)
            if(cy < 90 and up_taken== False):
                time_up = time.time()
                up_taken = True

            if(cy > 180 and down_taken == False):
                time_down = time.time()
                down_taken = True
        # print("Y pos is ", cy)
    if(time_down<time_up and down_taken and up_taken):
        print("Hand is moving UP")
        return 5
    if(time_up<time_down and down_taken and up_taken):
        print("Hand is moving DOWN")
        return 6
    return 0

def eulcidian_distance(X, Y, x_0, y_0):
    sumX = 0
    sumY = 0
    for x in X:
        sumX = sumX+(x-x_0)**2
    for y in Y:
        sumY = sumY+(y-y_0)**2
    return math.sqrt(sumX+sumY)

def utility(handLms, w, h):
    dist = 0
    X = []
    Y = []
    x_0 = 0
    y_0 = 0
    for id, lm in enumerate(handLms.landmark):
        #print(id,lm)
        # Height, width and channel of image
        h, w, c = img.shape
        # X and Y coordinate
        # their values in decimal so
        # we have to convert into pixel
        cx, cy = int((lm.x)*w),int((lm.y)*h)
        if id==0:
            x_0 = cx
            y_0 = cy
        elif (id%4)==0:
            X.append(cx)
            Y.append(cy)
    return (eulcidian_distance(X, Y, x_0, y_0))

# Create Video object
cap = cv2.VideoCapture(0)
# Formality we have to write before start
# using this model
mpHands = mp.solutions.hands
# Creating an object from class Hands
hands = mpHands.Hands(max_num_hands=1)
# creating an object to draw hand landmarks
mpDraw = mp.solutions.drawing_utils
# Previous time for frame rate
pTime = 0
# Current time for frame rate
cTime = 0
frame = 0
time_before_gesture = 0
time_vol = 0
while True:
    # Getting our Frame
    if((time.time() - time_before_gesture)>2.000):
        time_before_gesture = 0
    if((time.time() - time_vol)>0.1):
        time_vol = 0
    success, img = cap.read()
    img = cv2.resize(img,(500,500))
    # Convert image into RGB
    imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    # Calling the hands object to the getting results
    results = hands.process(imgRGB)
    #print(results.multi_hand_landmarks)
    # Checking something is detected or not
    if results.multi_hand_landmarks :
        # Extracting the multiple hands
        # Go through each hand
        for handLms in results.multi_hand_landmarks:
            frame = (frame+1)%31
            # Getting id(index number) and landmark of each hand
            for id, lm in enumerate(handLms.landmark):
                #print(id,lm)
                # Height, width and channel of image
                h, w, c = img.shape
                # X and Y coordinate
                # their values in decimal so
                # we have to convert into pixel
                cx, cy = int(lm.x*w),int(lm.y*h)
                #print(id, cx, cy)
                if id ==4:
                    cv2.circle(img, (cx, cy), 7, (255,0,255), cv2.FILLED)
            # Draw the landmarks and line of the each hands
            mpDraw.draw_landmarks(img, handLms, mpHands.HAND_CONNECTIONS)
            # openClose(frame, handLms)
            if(time_before_gesture == 0):
                hor = sliderHori(handLms, frame, w)
            if(time_vol == 0):
                ver = sliderVer(handLms, frame, h)
                # op = openClose(frame, handLms, w, h)
            if(hor>0):
                hor = 0
                time_before_gesture = time.time()
            if(ver>0):
                ver = 0
                time_vol = time.time()
            # if(op>0):
            #     op = 0
            #     time_before_gesture = time.time()
    else:
        left_taken = False
        right_taken = False
        up_taken = False
        down_taken = False
    # Getting the current time
    cTime = time.time()
    # Getting frame per second
    fps = 1/(cTime-pTime)
    # Previous time become current time
    pTime = cTime
    # Labeling the Frame rate
    cv2.putText(img,str(int(fps)),(10,70),cv2.FONT_HERSHEY_PLAIN,3,
                            (255,0,255),3)

    cv2.imshow("image",img) # Show the frame
    cv2.waitKey(1) # Wait for 1 millisecond


import logging
import random
import re

from Legobot.Lego import Lego
import requests

logger = logging.getLogger(__name__)

class XKCD(Lego):
    def listening_for(self, message):
        text = message.get('text')

        if text and text.startswith('!xkcd'):
            return True

        return False

    def handle(self, message):
        text = message['text']
        logger.debug(f'Handling message:\n{text}')
        comic_id = self._get_comic_id(text)
        response = self._get_comic(comic_id)

        if response:
            opts = self.build_reply_opts(message)
            self.reply(message, response, opts)

    def _call_xkcd_api(self, comic_id=None):
        if comic_id:
            url = f'https://xkcd.com/{comic_id}/info.0.json'
        else:
            url = 'https://xkcd.com/info.0.json'

        call = requests.get(url)

        if call.status_code == 200:
            return call.json()
        else:
            logger.error(
                f'Error on GET {url}: {call.status_code}: {call.text}')

            return None

    def _get_comic(self, comic_id):
        response = None
        comic = self._call_xkcd_api(comic_id)

        if comic:
            response = 'xkcd #{}: {} {}'.format(
                comic['num'],
                comic['alt'],
                comic['img']
            )

        return response

    def _get_comic_id(self, text):
        comic_id = None
        params = re.split(r'\s+', text)

        if len(params) > 1:
            comic_id = params[1]

        if comic_id in ('r', 'random'):
            comic_id = self._get_random_comic_id()

        return comic_id

    def _get_random_comic_id(self):
        comic_id = 1337
        latest = self._call_xkcd_api()

        if latest:
            latest_id = latest['num']
            comic_id = random.randint(1, latest_id)  # nosec

        return comic_id

    def get_name(self):
        return 'xkcd'

    def get_help(self):
        return 'Fetch an xkcd. Usage: !xkcd [r|random|<int>]'


import matplotlib.pyplot as plt
import tensorflow as tf

# 载入并准备好 MNIST 数据集。将样本从整数转换为浮点数
mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# 显示一张样本图片
plt.imshow(x_train[0])
plt.show()

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    # 数据集有 10 个类, Dense 层需要 10 个输出及一个 softmax 激活函数
    tf.keras.layers.Dense(10, activation='softmax')
])

tensorboard = tf.keras.callbacks.TensorBoard(log_dir='./log')

if False:
    # 顺序标记的label，使用sparse_categorical_crossentropy作为损失函数
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    model.fit(x_train, y_train, epochs=5, callbacks=[tensorboard])

    # 验证模型
    model.evaluate(x_test, y_test, verbose=2)
else:
    train_y_one_hot_label = tf.keras.utils.to_categorical(y_train)
    test_y_one_hot_label = tf.keras.utils.to_categorical(y_test)

    # 独热编码label，使用categorical_crossentropy作为损失函数
    model.compile(
        optimizer='adam',
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )

    model.fit(x_train, train_y_one_hot_label, epochs=5, callbacks=[tensorboard])
    # 验证模型
    model.evaluate(x_test, test_y_one_hot_label, verbose=2)


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.spectral_norm as spectral_norm

# This implementation directly uses SNGAN ResNet Generator and Discriminator

class ResNetBlock(nn.Module):
	def __init__(self, in_channel, out_channel, block_type = "generator", blockno=0):
		super(ResNetBlock, self).__init__()
		self.bn1 = nn.BatchNorm2d(in_channel)
		self.act1 = nn.ReLU()
		self.conv1 = nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1) if block_type == "generator" else spectral_norm(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1))
		self.bn2 = nn.BatchNorm2d(out_channel)
		self.act2 = nn.ReLU()
		self.conv2 = nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1) if block_type == "generator" else spectral_norm(nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1))
		self.conv_residual = nn.Conv2d(in_channel, out_channel, kernel_size=1, padding=0) if block_type == "generator" else spectral_norm(nn.Conv2d(in_channel, out_channel, kernel_size=1, padding=0))
		self.block_type = block_type
		self.blockno = blockno

	def forward(self, x):
		b_size, channel, h, w = x.size()
		if self.block_type == "generator":
			x = self.bn1(x)
		out = self.act1(x)
		out = self.conv1(out)
		if self.block_type == "generator":
			out = F.interpolate(out, size=(h*2,w*2), mode="bilinear")
			out = self.bn2(out)
		out = self.act2(out)
		out = self.conv2(out)
		if (self.block_type == "discriminator") and (self.blockno == 0 or self.blockno == 1):
			out = F.avg_pool2d(out, 2)

		residual = self.conv_residual(x)
		if self.block_type == "generator":
			residual = F.interpolate(residual, size=(h*2,w*2), mode="bilinear")

		elif (self.block_type == "discriminator") and (self.blockno == 0 or self.blockno == 1):
			residual = F.avg_pool2d(residual, 2)

		return residual + out

class SNGAN_Generator(nn.Module):
	def __init__(self):
		super(SNGAN_Generator, self).__init__()
		self.linear = nn.Linear(128, 4*4*256) # z = [batch, 128] -> [batch, 256, 4, 4]
		resblocks = []
		resblocks += [ResNetBlock(256, 256, block_type="generator")]
		resblocks += [ResNetBlock(256, 256, block_type="generator")]
		resblocks += [ResNetBlock(256, 256, block_type="generator")]
		self.resblocks = nn.Sequential(*resblocks)
		self.bn = nn.BatchNorm2d(256)
		self.act = nn.ReLU()
		self.conv = nn.Conv2d(256, 3, kernel_size=3, padding=1)
		self.tanh = nn.Tanh()

	def forward(self, z):
		b_size, z_dim = z.size()
		out = self.linear(z)
		out = out.view(b_size, 256, 4, 4)
		out = self.resblocks(out)
		out = self.conv(self.act(self.bn(out)))
		return self.tanh(out)

class SNGAN_Discriminator(nn.Module):
	def __init__(self):
		super(SNGAN_Discriminator, self).__init__()
		resblocks = []
		resblocks += [ResNetBlock(3, 128, block_type="discriminator", blockno=0)]
		resblocks += [ResNetBlock(128, 128, block_type="discriminator", blockno=1)]
		resblocks += [ResNetBlock(128, 128, block_type="discriminator", blockno=2)]
		resblocks += [ResNetBlock(128, 128, block_type="discriminator", blockno=3)]
		self.resblocks = nn.Sequential(*resblocks)
		self.act = nn.ReLU()
		self.linear = spectral_norm(nn.Linear(128, 1))

	def forward(self, x):
		out = self.resblocks(x)
		out = self.act(out)
		out = F.avg_pool2d(out, out.size(2))
		out = out.view(out.size(0), -1)
		return self.linear(out)


"""initial listing model

Revision ID: d6c2334f4816
Revises:
Create Date: 2021-09-27 17:33:02.166128

"""
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

from alembic import op

# revision identifiers, used by Alembic.
revision = "d6c2334f4816"
down_revision = None
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "catalog_listing",
        sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False),
        sa.Column("data", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
        sa.PrimaryKeyConstraint("id"),
    )
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table("catalog_listing")
    # ### end Alembic commands ###


#coding=utf-8

def read_txt(path,match,splitSymbol,lr=1):
    with open(path, 'r') as f:
        for line in f.readlines():
            if match in line:
                value = line.split(splitSymbol)[lr]
                return value.strip('\n')

if __name__ == "__main__":
    from public.readConfig import *
    headerfilepath = headerManage['head_path']
    match = headerManage['init_cookie']
    joinSymbol = "==>"
    print(read_txt(headerfilepath,match,joinSymbol,1))


# -*- coding: utf-8 -*-

import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
# from keras.callbacks import Callback
from sklearn.metrics import confusion_matrix, f1_score, precision_score, recall_score

import tokenization

import tensorflow as tf
import tensorflow_hub as hub
from sklearn import preprocessing
from sklearn.model_selection import train_test_split
from sklearn.model_selection import KFold
from sklearn.utils import shuffle
from sklearn.metrics import f1_score

# Importing argparse for command line arguments.
import argparse

# Parsing command line arguments.
parser = argparse.ArgumentParser(description='Train the proposed model. An example on how to run the script is as follows: \n python main.py --model=comemnet-bilstm')
parser.add_argument('-m','--model',default='comemnet-bilstm',  help="Can be either 'comemnet-lstm', 'comemnet-bilstm' or 'erin'.\n Example: model=comemnet-bilstm (default=comemnet-bilstm)")
parser.add_argument('-d','--demo', default=False, help="Boolean. Whether to create a gradio demo (default=False).")
parser.add_argument('-p', '--pretrained', default=False, help='Boolean. Whether to use pretrained model. (default=False). If pretrained=True, please download the google drive files mentioned in the README.')

args = vars(parser.parse_args())
print("You selected the following parameters for the script to run.")
print(args)

if args['demo']:
    import gradio as gr

# print("Using TensorFlow version: ",tf.__version__)

#checking for GPU device, otherwise model will be trained on CPU
device_name = tf.test.gpu_device_name()
print(device_name)
if device_name != '/device:GPU:0':
  print('GPU device not found, falling back to CPU. Warning, code will run slowly on CPU.')
print('Found GPU at: {}'.format(device_name))

#data is initially parsed from JSON files
#now preprocesssing the parsed data for model training
print("Loading and preprocessing data")
data_path = 'Processed Data/Input_500_29_relation.tsv'

train_data = pd.read_csv(data_path, encoding='utf-8', sep = '\t')

#filling nan values with empty strings
train_data.fillna("", inplace = True)

# Shuffle data so that there is a higher chance of the train and test data being from the same distribution.
train_data = shuffle(train_data, random_state = 1)

# # Now read the rows, convert them into strings and then only keep the unique ones.
sentences_and_labels =  np.array([[' '.join(map(str, row[:-1].tolist())).strip(), row[-1]] for row in train_data.iloc[:,:].values])

sentences = sentences_and_labels[:, 0]

#using label encoder to map relations to integer lablels
label = preprocessing.LabelEncoder()
y = label.fit_transform(train_data['relation'])
label_mappings = integer_mapping = {i: l for i, l in enumerate(label.classes_)}

def get_labels(y_pred):
    """
    This method generates integer labels from one-hot representations;
    Returns integer label values

    y_pred -- specifies one-hot labels generated by our model
    """
    y_pred_label = np.zeros((len(y_pred), 1))
    print(y_pred_label.shape)
    for index in range(len(y_pred)):
        y_pred_label[index] = np.argmax(y_pred[index])
    return y_pred_label

#if pretrained is set to False, then initialize a new model and train it
if args['pretrained'] == False:
    print("Retraining model from scratch.")
    print("Loading bert from TensorFlow-Hub, this may take a while... Please be patient.")
    m_url = 'https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/2'
    bert_layer = hub.KerasLayer(m_url, trainable=False)

    vocab_file = bert_layer.resolved_object.vocab_file.asset_path.numpy()
    do_lower_case = bert_layer.resolved_object.do_lower_case.numpy()
    tokenizer = tokenization.FullTokenizer(vocab_file, do_lower_case)

    def bert_encode(texts, tokenizer, max_len=512):
        all_tokens = []
        all_masks = []
        all_segments = []

        for text in texts:
            text = tokenizer.tokenize(text)

            text = text[:max_len-2]
            input_sequence = ["[CLS]"] + text + ["[SEP]"]
            pad_len = max_len-len(input_sequence)

            tokens = tokenizer.convert_tokens_to_ids(input_sequence) + [0] * pad_len
            pad_masks = [1] * len(input_sequence) + [0] * pad_len
            segment_ids = [0] * max_len

            all_tokens.append(tokens)
            all_masks.append(pad_masks)
            all_segments.append(segment_ids)

        return np.array(all_tokens), np.array(all_masks), np.array(all_segments)

    def build_model_erin(bert_layer, max_len=512):
        """
        This method reimplements Macdonald and Barbosa's model;
        Use the reported parameters from their paper to ensure fair comparison;
        Returns model object

        bert_layer -- specifies the pretrained BERT layer for generating embedding vectors
        max_len --  specifies maximum length of input tokens for BERT layer
        """
        input_word_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_word_ids")
        input_mask = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_mask")
        segment_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="segment_ids")

        pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])

        #print(tf.shape(sequence_output))
        clf_output = sequence_output[:, :, :]
        #print(tf.shape(clf_output))

        #build model with one LSTM layer, followed by fully connected and softmax layers
        lay = tf.keras.layers.LSTM(1, return_sequences=True)(clf_output)
        lay = tf.keras.layers.Flatten()(lay)
        out = tf.keras.layers.Dense(29, activation='softmax')(lay)

        model = tf.keras.models.Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=out)
        #utilize RMSProp optimizer as reported by Macdonald and Barbosa
        model.compile(tf.keras.optimizers.RMSprop(lr=0.001), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

        return model

    def build_model_comemnet_lstm(bert_layer, max_len=512):
        """
        This method implements our baseline model;
        baseline model comprises of a CNN layer and an LSTM layer;
        Returns model object

        bert_layer -- specifies the pretrained BERT layer for generating embedding vectors
        max_len --  specifies maximum length of input tokens for BERT layer
        """
        input_word_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_word_ids")
        input_mask = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_mask")
        segment_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="segment_ids")

        pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])
        # np.savez_compressed(f"bert_sequence_op_seed_{seed}.npz", sequence_output)

        #print(tf.shape(sequence_output))
        clf_output = sequence_output[:, :, :]
        #print(tf.shape(clf_output))

        #build model with one CNN layer with 8 filters and one LSTM layer with 8 hidden units
        lay = tf.keras.layers.Conv1D(filters=8, kernel_size=5, strides=1, padding="same", activation="relu")(clf_output)
        lay = tf.keras.layers.MaxPooling1D(2, 2)(lay)
        lay = tf.keras.layers.LSTM(8, return_sequences=True, dropout=0.2)(lay)
        lay = tf.keras.layers.Flatten()(lay)
        out = tf.keras.layers.Dense(29, activation='softmax')(lay)

        model = tf.keras.models.Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=out)
        #use sparse categorical crossentropy as output labels are defined as integer values
        model.compile(tf.keras.optimizers.Adam(lr=2e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

        return model

    def build_model_comemnet_bilstm(bert_layer, max_len=512):
        """
        This method is an improvement over our baseline;
        improved model includes a CNN and BiLSTM network;
        Returns model object

        bert_layer -- specifies the pretrained BERT layer for generating embedding vectors
        max_len --  specifies maximum length of input tokens for BERT layer
        """
        input_word_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_word_ids")
        input_mask = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_mask")
        segment_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="segment_ids")

        pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])
        # np.savez_compressed(f"bert_sequence_op_seed_{seed}.npz", sequence_output)

        # print(tf.shape(sequence_output))
        clf_output = sequence_output[:, :, :]
        # print(tf.shape(clf_output))

        #improve our baseline model by adding BiLSTM layer for modeling dependencies in both directions
        lay = tf.keras.layers.Conv1D(filters=8, kernel_size=5, strides=1, padding="same", activation="relu")(clf_output)
        lay = tf.keras.layers.MaxPooling1D(2, 2)(lay)
        #implement dropout to prevent overfitting and ensure generalizability
        lay = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(8, return_sequences=True, dropout=0.2))(lay)
        lay = tf.keras.layers.Flatten()(lay)
        out = tf.keras.layers.Dense(29, activation='softmax')(lay)

        model = tf.keras.models.Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=out)
        model.compile(tf.keras.optimizers.Adam(lr=2e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

        return model

    """### Obtaining Train, test splits.
    ###### In the train splits, we will have a separate validation split.
    """

    #specifying checkpoint directory so trained model is saved after each epoch
    print("Checkpoint will only be saved for the last epoch.")
    checkpoint_path = "training_relations/cp.ckpt"
    checkpoint_dir = os.path.dirname(checkpoint_path)

    """# Do not run the following cell if using checkpointed files."""

    """
    Perform evaluation using Monte-carlo CV;
    model was tested on 5 different seeds and all metrics were averaged across seeds for reporting final results;
    model hyperparameters was tuned using 5 fold CV on development set (20% of training set)
    """

    with tf.device(device_name):
        splits = 5 # For five fold cross-validation.
        seeds = [i for i in range(splits)]  # Fix the seed value for reproducibility.
        # seeds = [0]

        val_dict = {}
        test_dict = {}
        test_accs = []
        test_f1s = []

        # First get random train-test splits. Doesn't include validation, which will be obtained from the train set.
        for seed in seeds:
            print(f"Training model for seed {seed}.")
            x_t, x_test, y_t, y_test = train_test_split(sentences, y, random_state=seed, test_size=0.2)   # Global training and test sets.

            # Now get validation sets from each training set.
            # kf = KFold(n_splits=5, shuffle=False) # Setting shuffle=False because shuffled dataset already before.
            # fold_count = 0

            # for train_index, val_index in kf.split(x_t):
            #     #print(x_t.shape)
            #     #print(y_t.shape)
            #     x_train, x_val = x_t[train_index], x_t[val_index]   # Training and validation features.
            #     y_train, y_val = y_t[train_index], y_t[val_index]   # Training and validation labels.

            #     #encode train data
            #     max_len = 80
            #     train_input = bert_encode(x_train, tokenizer, max_len=max_len)
            #     train_labels = y_train
            #     x_val = bert_encode(x_val, tokenizer, max_len=max_len)

            #     model = build_model(bert_layer, max_len=max_len)
            #     model.summary()
            #     #checkpoint = tf.keras.callbacks.ModelCheckpoint('model.h5', monitor='val_accuracy', save_best_only=True, verbose=1)
            #     checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path, save_weights_only=True, verbose=1)
            #     earlystopping = tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', patience=5, verbose=1)
            #     train_sh = model.fit(
            #     train_input, train_labels,
            #     #validation_split=0.2,
            #     validation_data=(x_val, y_val),
            #     epochs=2,
            #     callbacks=[checkpoint, earlystopping],
            #     batch_size=16,
            #     verbose=1)

            #     # Validation sets can be used for hyperparamter tuning.
            #     val_dict[str(seed) + str(fold_count)] = train_sh.history
            #     fold_count += 1

            #encode whole train data
            #specify max length parameters for our models
            max_len = 80
            #Macdonald and Barbosa considered 50 max length for their proposed model
            #we preserve their reported hyperparameters for replicating their results
            erin_max_len = 50

            print("Encoding input through BERT encoder.")
            if args['model'] == 'erin':
                train_input = bert_encode(x_t, tokenizer, max_len=erin_max_len)
                train_labels = y_t
            else:
                train_input = bert_encode(x_t, tokenizer, max_len=max_len)
                train_labels = y_t

            #train model on whole train data
            if args['model'] == "comemnet-bilstm":
                model = build_model_comemnet_bilstm(bert_layer, max_len=max_len)
            elif args['model'] =='comemnet-lstm':
                model = build_model_comemnet_lstm(bert_layer, max_len=max_len)
            elif args['model'] == 'erin':
                model = build_model_erin(bert_layer, max_len=erin_max_len)

            model.summary()

            #checkpoint = tf.keras.callbacks.ModelCheckpoint('model.h5', monitor='val_accuracy', save_best_only=True, verbose=1)
            checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path, save_weights_only=True, verbose=1)
            earlystopping = tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', patience=5, verbose=1)
            train_sh = model.fit(
            train_input, train_labels,
            validation_split=0,
            epochs=6,
            callbacks=[checkpoint, earlystopping],
            batch_size=16,
            verbose=1)

            #encode test data
            if args['model'] == 'erin':
                test_input = bert_encode(x_test, tokenizer, max_len=erin_max_len)
            else:
                test_input = bert_encode(x_test, tokenizer, max_len=max_len)

            # Evaluate the model on the test data using `evaluate`
            print("Evaluating on test data for ", seed)
            results = model.evaluate(test_input, y_test, batch_size=16)

            #calculate F1-score
            y_pred = model.predict(test_input, verbose=1)
            y_pred_label = get_labels(y_pred)
            f1_value = f1_score(y_test, y_pred_label, average='macro')
            results.append(f1_value)
            test_f1s.append(f1_value)
            test_accs.append(results[1])
            print(f"Test loss, Test acc, F1-score for seed {seed}: ", results)

            test_dict[seed] = results
        print("Average accuracy: ", np.mean(test_accs))
        print("Average F1 score: ", np.mean(test_f1s))
        print("Results for all seed values: ", test_dict)

"""## The following section is for loading the checkpoint weights and then obtaining the confusion matrix."""

if args['pretrained']:
    print("Currently confusion matrix is only supported for CoMeMNet-BiLSTM")
    cms = []
    # #code for loading weights from checkpoint
    # with tf.device(device_name):
    #     results = []
    #     splits = 5 # For five fold cross-validation.
    #     #seeds = [i for i in range(splits)]  # Fix the seed value for reproducibility.
    #     seeds = [4]#0, 1, 2, 3, 4]
    #
    #     val_dict = {}
    #     test_dict = {}
    #
    #     # First get random train-test splits. Doesn't include validation, which will be obtained from the train set.
    #     for seed in seeds:
    #         x_t, x_test, y_t, y_test = train_test_split(sentences, y, random_state=seed, test_size=0.2)   # Global training and test sets.
    #
    #         # Evaluate the new model
    #         max_len = 80
    #         if args['model'] == "comemnet-bilstm":
    #             model = build_model_comemnet_bilstm(bert_layer, max_len=max_len)
    #         elif args['model'] == 'comemnet-lstm':
    #             model = build_model_comemnet_lstm(bert_layer, max_len=max_len)
    #         elif args['model'] == 'erin':
    #             model = build_model_erin(bert_layer, max_len=max_len)
    #         model.summary()
    #         # Loads the weights for new model
    #         checkpoint_path = f"/content/drive/Shareddrives/CMPUT 656 Data and Results/Result 29 Relations/Seed {seed}/training_relations/cp.ckpt"
    #         training_relations_path = f"/content/drive/Shareddrives/CMPUT 656 Data and Results/Result 29 Relations/Seed {seed}/training_relations/"
    #         # /content/drive/Shareddrives/CMPUT 656 Data and Results/Result 29 Relations/Seed 0/training_relations/cp.ckpt.data-00000-of-00001
    #
    #
    #         model.load_weights(checkpoint_path)
    #
    #         #encode test data
    #         test_input = bert_encode(x_test, tokenizer, max_len=max_len)
    #
    #         # Evaluate the model on the test data using `evaluate`
    #         print("Evaluate on test data for ", seed)
    #         # results = new_model.evaluate(test_input, y_test, batch_size=16)
    #
    #         #calculate F1-score
    #         y_pred = model.predict(test_input, batch_size=16, verbose=1)
    #         y_pred_label = get_labels(y_pred)
    #         f1_value = f1_score(y_test, y_pred_label, average='macro')
    #         np.savez_compressed(training_relations_path +  "predictions.npz", y_pred_label)
    #         cm = confusion_matrix(y_pred_label, y_test, labels=[i for i in label_mappings.keys()])
    #         cms.append(cm)
    #
    #         results.append(f1_value)
    #         print("Test loss, Test acc, F1-score:", results)

    for k, v in label_mappings.items():
        print(k, v)

    #mapping each relations to unique integer labels
    reduced_label_mappings = {
        0: 'None',
        1: 'award=nominee',
        2: 'author-works_written',
        3: 'book-genre',
        4: 'company-industry',
        5: 'person-graduate',
        6: 'actor-character',
        7: 'director-film',
        8: 'film-country',
        9: 'film-genre',
        10: 'film-language',
        11: 'film-music',
        12: 'film-production_company',
        13: 'actor-film',
        14: 'producer-film',
        15: 'writer-film',
        16: 'political_party-politician',
        17: 'location-contains',
        18: 'musician-album',
        19: 'musician-origin',
        20: 'person-place_of_death',
        21: 'person-nationality',
        22: 'person-parents',
        23: 'person-place_of_birth',
        24: 'person-profession',
        25: 'person-religion',
        26: 'person-spouse',
        27: 'football_position-player',
        28: 'sports_team-player'
    }

    """# Read in all the predictions for different seeds and make a confusion matrix by averaging them together.

    NOTE: The confusion matrix is normalized.
    """

    cms = []
    seeds = [0, 1, 2, 3, 4]
    for seed in seeds:
        training_relations_path = f"Result 29 Relations/Seed {seed}/training_relations/"
        x_t, x_test, y_t, y_test = train_test_split(sentences, y, random_state=seed, test_size=0.2)
        del x_t
        del x_test
        del y_t
        predictions = np.load(training_relations_path + "predictions.npz", allow_pickle=True)['arr_0'].tolist()
        cms.append(confusion_matrix(predictions, y_test, labels=[i for i in label_mappings.keys()]))

    plt.clf()
    averaged_cms = np.mean(cms, axis=0)
    averaged_cms = averaged_cms.astype('float') / averaged_cms.sum(axis=1)[:, np.newaxis]
    cms_df = pd.DataFrame(averaged_cms, index = [value for value in reduced_label_mappings.values()],
                          columns=[value for value in reduced_label_mappings.values()])
    fig, ax = plt.subplots(figsize=(15, 10))
    heat = sns.heatmap(cms_df, vmin=0.0, vmax=1.0, cbar_kws={'label': 'Normalized value'})
    heat.figure.axes[-1].yaxis.label.set_size(20)
    yticks = [i.upper() for i in cms_df.index]
    xticks = [i.upper() for i in cms_df.columns]
    plt.yticks(plt.yticks()[0], labels=yticks, rotation=0)
    plt.xticks(plt.xticks()[0], labels=xticks, rotation=270)
    plt.title("Confusion matrix for all 29 relations - CoMemNet-BiLSTM", fontsize=20)
    plt.show()

"""# Gradio app to demo the model."""

if args['demo']:
    print("Loading BERT encoder, this might take some time... Please be patient.")
    m_url = 'https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/2'
    bert_layer = hub.KerasLayer(m_url, trainable=False)

    def get_labels(y_pred):
        y_pred_label = np.zeros((len(y_pred),1))
        print(y_pred_label.shape)
        for index in range(len(y_pred)):
            y_pred_label[index] = np.argmax(y_pred[index])
        return y_pred_label

    vocab_file = bert_layer.resolved_object.vocab_file.asset_path.numpy()
    do_lower_case = bert_layer.resolved_object.do_lower_case.numpy()
    tokenizer = tokenization.FullTokenizer(vocab_file, do_lower_case)

    def bert_encode(texts, tokenizer, max_len=512):
        all_tokens = []
        all_masks = []
        all_segments = []

        for text in texts:
            text = tokenizer.tokenize(text)

            text = text[:max_len-2]
            input_sequence = ["[CLS]"] + text + ["[SEP]"]
            pad_len = max_len-len(input_sequence)

            tokens = tokenizer.convert_tokens_to_ids(input_sequence) + [0] * pad_len
            pad_masks = [1] * len(input_sequence) + [0] * pad_len
            segment_ids = [0] * max_len

            all_tokens.append(tokens)
            all_masks.append(pad_masks)
            all_segments.append(segment_ids)

        return np.array(all_tokens), np.array(all_masks), np.array(all_segments)

    def build_model(bert_layer, max_len=512):
        input_word_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_word_ids")
        input_mask = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_mask")
        segment_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="segment_ids")

        pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])

        # print(tf.shape(sequence_output))
        clf_output = sequence_output[:, :, :]
        # print(tf.shape(clf_output))

        lay = tf.keras.layers.Conv1D(filters=8, kernel_size=5, strides=1, padding="same", activation="relu")(clf_output)
        lay = tf.keras.layers.MaxPooling1D(2, 2)(lay)
        lay = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(8, return_sequences=True, dropout=0.2))(lay)
        lay = tf.keras.layers.Flatten()(lay)
        out = tf.keras.layers.Dense(29, activation='softmax')(lay)

        model = tf.keras.models.Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=out)
        model.compile(tf.keras.optimizers.Adam(lr=2e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

        return model

    reduced_label_mappings = {
        0: 'None',
        1: 'award-nominee',
        2: 'author-works_written',
        3: 'book-genre',
        4: 'company-industry',
        5: 'person-graduate',
        6: 'actor-character',
        7: 'director-film',
        8: 'film-country',
        9: 'film-genre',
        10: 'film-language',
        11: 'film-music',
        12: 'film-production_company',
        13: 'actor-film',
        14: 'producer-film',
        15: 'writer-film',
        16: 'political_party-politician',
        17: 'location-contains',
        18: 'musician-album',
        19: 'musician-origin',
        20: 'person-place_of_death',
        21: 'person-nationality',
        22: 'person-parents',
        23: 'person-place_of_birth',
        24: 'person-profession',
        25: 'person-religion',
        26: 'person-spouse',
        27: 'football_position-player',
        28: 'sports_team-player'
    }

    seed = 0
    print(f"Using model checkpoint file for seed {seed}")
    checkpoint_path = f"Result 29 Relations/Seed {seed}/training_relations/cp.ckpt"
    max_len = 80

    new_model = build_model(bert_layer, max_len=max_len)
    new_model.load_weights(checkpoint_path)

    def predict_relation(sentence):
        print("Encoding input through BERT encoder.")
        test_input = bert_encode(np.array([sentence]), tokenizer, max_len=max_len)
        y_pred = new_model.predict(test_input, batch_size=16, verbose=1)
        y_pred_label = get_labels(y_pred)
        relations = [rel for rel in reduced_label_mappings.values()]
        probabilities = y_pred.tolist()[0]
        result_dict = {}
        for k, v in zip(relations, probabilities):
            result_dict[k] = v
        return result_dict

    iface = gr.Interface(
        predict_relation,
        inputs="text",
        outputs="label",
        interpretation="default",
        title="CoMemNet - Relation Extractor", description="NOTE: Model is trained on Wikipedia table data and not continuous text."
    )

    iface.launch(debug=False)


import torch
import torch.nn as nn
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
device = torch.device("cuda:0")

com = [
    16,
    "ELU",
    32,
    "ELU",
    64,
    "ELU",
    128,
    "ELU",
    256,
    "ELU",
    128,
    "ELU",
    64,
    "ELU",
    32,
    "ELU",
    4,
]
rec = [
    "SIGMOID",
    32,
    "ELU",
    64,
    "ELU",
    128,
    "ELU",
    256,
    "ELU",
    128,
    "ELU",
    64,
    "ELU",
    32,
    "ELU",
    16,
    "ELU",
    3,
]

class ComDefend(nn.Module):
    def __init__(self, phi=20):
        super(ComDefend, self).__init__()
        self.name = "ComDefend"
        self.phi = phi
        self.com, self.rec = self._make_layers(com, rec)

    def _make_layers(self, com, rec):
        comCNN = []
        recCNN = []
        inchannels = 3

        for x in com:
            if x == "ELU":
                comCNN.append(nn.ELU())
            elif x == "SIGMOID":
                comCNN.append(nn.Sigmoid())
            else:
                comCNN.append(nn.Conv2d(inchannels, x, kernel_size=3, padding=1))
                comCNN.append(nn.BatchNorm2d(x))
                inchannels = x

        for x in rec:
            if x == "ELU":
                recCNN.append(nn.ELU())
            elif x == "SIGMOID":
                recCNN.append(nn.Sigmoid())
            else:
                recCNN.append(nn.Conv2d(inchannels, x, kernel_size=3, padding=1))
                recCNN.append(nn.BatchNorm2d(x))
                inchannels = x

        return nn.Sequential(*comCNN), nn.Sequential(*recCNN)

    def forward(self, x):
        hidden = self.com(x)
        hidden_gauss = hidden + (self.phi ** 0.5) * torch.randn(hidden.shape).to(device)
        output = self.rec(hidden_gauss)

        return hidden, output

comDefend = ComDefend().to(device)

import numpy as np

data = (
    np.load("/home/liyanni/1307/zwh/whiteBox/data/cifar10_data.npy").astype("float32")
    / 255
)
data = data.transpose((0, 3, 1, 2))
label = np.load("/home/liyanni/1307/zwh/whiteBox/data/cifar10_label.npy")
label = np.reshape(label, (label.shape[0],))

def dataset(batch=50):
    train = data[:50000]
    for i in range(len(train) // batch):
        yield train[i * batch : (i + 1) * batch], train[i * batch : (i + 1) * batch]

comDefend.train()
from torch.autograd import Variable

criterion1 = nn.MSELoss()
criterion2 = nn.MSELoss()
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, comDefend.parameters()))
for epoch in range(20):
    dataGen = dataset()
    l1 = 0.0
    l2 = 0.0
    for i, (x, y) in enumerate(dataGen, 1):
        # get the inputs
        x, y = torch.from_numpy(x), torch.from_numpy(y)
        x, y = x.to(device), y.to(device)
        x, y = Variable(x, requires_grad=True), Variable(y)
        # zero the parameter gradients
        optimizer.zero_grad()
        # forward
        o1, o2 = comDefend(x)
        # loss
        loss1 = criterion1(o1, torch.zeros(o1.shape).to(device))
        loss2 = criterion2(o2, y)
        loss = 0.0001 * loss1 + loss2
        # backward
        loss.backward()
        # update weights
        optimizer.step()
        # print statistics
        l1 += loss1.item()
        l2 += loss2.item()

        if i % 50 == 0:  # print every 2000 mini-batches
            print("epoch:{}  loss1:{}  loss2{}".format(epoch, l1 / 50.0, l2 / 50.0))
            l1, l2 = 0.0, 0.0

from models.resnet18 import ResNet50
from collections import OrderedDict

target = ResNet50()
pthfile = r"checkpoints/resnet50_ckpt.pth"
d = torch.load(pthfile)["net"]
d = OrderedDict([(k[7:], v) for (k, v) in d.items()])
target.load_state_dict(d)
target.to(device)

mean = np.mean(data, axis=0)

adv_x = (
    np.load("../adv_data/cifar10/cw/0.015x_adv.npy")
    .astype("float32")
    .transpose(0, 3, 1, 2)
)
adv_x = adv_x + mean

def testGen(data, batch=50):
    testx = data[50000:]
    testy = label[50000:]
    for i in range(len(testx) // batch):
        yield testx[i * batch : (i + 1) * batch], testy[i * batch : (i + 1) * batch]

test = testGen(adv_x)
correct = 0
total = 0
for inputs, labels in test:
    inputs, labels = torch.from_numpy(inputs), torch.from_numpy(labels)
    inputs, labels = inputs.to(device), labels.to(device)
    _, outputs = comDefend(inputs)
    outputs = target(outputs)
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum()
print("Accuracy {}".format(100.0 * correct / total))


import os
import json

class fiveThreeNineFeatureUtil(object):
    logPath = os.path.join('Log', 'addInformation')

    lockCountsDict = {
        "0": {},
        "1": {},
        "2": {},
        "3": {},
        "4": {},
        "5": {},
        "6": {},
        "7": {},
        "8": {},
        "9": {}
    }

    currentLockDict = {
        "0": 0,
        "1": 0,
        "2": 0,
        "3": 0,
        "4": 0,
        "5": 0,
        "6": 0,
        "7": 0,
        "8": 0,
        "9": 0
    }

    def __init__(self):
        if os.path.exists(self.logPath):
            os.remove(self.logPath)

    def addInformation(self, tailNumbers, issue):
        for key in self.currentLockDict:
            if key in tailNumbers:
                if self.currentLockDict[key] not in self.lockCountsDict[key]:
                    self.lockCountsDict[key].update({
                        self.currentLockDict[key]: 1
                    })
                else:
                    self.lockCountsDict[key][self.currentLockDict[key]] += 1
                self.currentLockDict[key] = 0
            else:
                self.currentLockDict[key] += 1

        with open(self.logPath, 'a+') as logFile:
            logFile.write('Issue: %s\n' % issue)
            logFile.write('lockCountsDict:\n%s\n' % json.dumps(self.lockCountsDict, indent=4))
            logFile.write('currentLockDict:\n%s\n' % json.dumps(self.currentLockDict, indent=4))
            logFile.write('======================\n')