import torch
from torch import nn, einsum
import numpy as np
from einops import rearrange, repeat

class CyclicShift(nn.Module):
    def __init__(self, displacement):
        super().__init__()
        self.displacement = displacement

    def forward(self, x):
        return torch.roll(x, shifts=(self.displacement, self.displacement), dims=(1, 2))

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, dim),
        )

    def forward(self, x):
        return self.net(x)

def create_mask(window_size, displacement, upper_lower, left_right):
    mask = torch.zeros(window_size ** 2, window_size ** 2)

    if upper_lower:
        mask[-displacement * window_size:, :-displacement * window_size] = float('-inf')
        mask[:-displacement * window_size, -displacement * window_size:] = float('-inf')

    if left_right:
        mask = rearrange(mask, '(h1 w1) (h2 w2) -> h1 w1 h2 w2', h1=window_size, h2=window_size)
        mask[:, -displacement:, :, :-displacement] = float('-inf')
        mask[:, :-displacement, :, -displacement:] = float('-inf')
        mask = rearrange(mask, 'h1 w1 h2 w2 -> (h1 w1) (h2 w2)')

    return mask

def get_relative_distances(window_size):
    indices = torch.tensor(np.array([[x, y] for x in range(window_size) for y in range(window_size)]))
    distances = indices[None, :, :] - indices[:, None, :]
    return distances

class WindowAttention(nn.Module):
    def __init__(self, dim, heads, head_dim, shifted, window_size, relative_pos_embedding):
        super().__init__()
        inner_dim = head_dim * heads

        self.heads = heads
        self.scale = head_dim ** -0.5
        self.window_size = window_size
        self.relative_pos_embedding = relative_pos_embedding
        self.shifted = shifted

        if self.shifted:
            displacement = window_size // 2
            self.cyclic_shift = CyclicShift(-displacement)
            self.cyclic_back_shift = CyclicShift(displacement)
            self.upper_lower_mask = nn.Parameter(create_mask(window_size=window_size, displacement=displacement,
                                                             upper_lower=True, left_right=False), requires_grad=False)
            self.left_right_mask = nn.Parameter(create_mask(window_size=window_size, displacement=displacement,
                                                            upper_lower=False, left_right=True), requires_grad=False)

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)

        if self.relative_pos_embedding:
            self.relative_indices = get_relative_distances(window_size) + window_size - 1
            self.pos_embedding = nn.Parameter(torch.randn(2 * window_size - 1, 2 * window_size - 1))
        else:
            self.pos_embedding = nn.Parameter(torch.randn(window_size ** 2, window_size ** 2))

        self.to_out = nn.Linear(inner_dim, dim)

    def forward(self, x):
        if self.shifted:
            x = self.cyclic_shift(x)

        b, n_h, n_w, _, h = *x.shape, self.heads

        qkv = self.to_qkv(x).chunk(3, dim=-1)
        nw_h = n_h // self.window_size
        nw_w = n_w // self.window_size

        q, k, v = map(
            lambda t: rearrange(t, 'b (nw_h w_h) (nw_w w_w) (h d) -> b h (nw_h nw_w) (w_h w_w) d',
                                h=h, w_h=self.window_size, w_w=self.window_size), qkv)

        dots = einsum('b h w i d, b h w j d -> b h w i j', q, k) * self.scale

        if self.relative_pos_embedding:
            dots += self.pos_embedding[self.relative_indices[:, :, 0], self.relative_indices[:, :, 1]]
        else:
            dots += self.pos_embedding

        if self.shifted:
            dots[:, :, -nw_w:] += self.upper_lower_mask
            dots[:, :, nw_w - 1::nw_w] += self.left_right_mask

        attn = dots.softmax(dim=-1)

        out = einsum('b h w i j, b h w j d -> b h w i d', attn, v)
        out = rearrange(out, 'b h (nw_h nw_w) (w_h w_w) d -> b (nw_h w_h) (nw_w w_w) (h d)',
                        h=h, w_h=self.window_size, w_w=self.window_size, nw_h=nw_h, nw_w=nw_w)
        out = self.to_out(out)

        if self.shifted:
            out = self.cyclic_back_shift(out)
        return out

class SwinBlock(nn.Module):
    def __init__(self, dim, heads, head_dim, mlp_dim, shifted, window_size, relative_pos_embedding):
        super().__init__()
        self.attention_block = Residual(PreNorm(dim, WindowAttention(dim=dim,
                                                                     heads=heads,
                                                                     head_dim=head_dim,
                                                                     shifted=shifted,
                                                                     window_size=window_size,
                                                                     relative_pos_embedding=relative_pos_embedding)))
        self.mlp_block = Residual(PreNorm(dim, FeedForward(dim=dim, hidden_dim=mlp_dim)))

    def forward(self, x):
        x = self.attention_block(x)
        x = self.mlp_block(x)
        return x

class PatchMerging(nn.Module):
    def __init__(self, in_channels, out_channels, downscaling_factor):
        super().__init__()
        self.downscaling_factor = downscaling_factor
        self.patch_merge = nn.Unfold(kernel_size=downscaling_factor, stride=downscaling_factor, padding=0)
        self.linear = nn.Linear(in_channels * downscaling_factor ** 2, out_channels)

    def forward(self, x):
        b, c, h, w = x.shape
        new_h, new_w = h // self.downscaling_factor, w // self.downscaling_factor
        x = self.patch_merge(x).view(b, -1, new_h, new_w).permute(0, 2, 3, 1)
        x = self.linear(x)
        return x

class StageModule(nn.Module):
    def __init__(self, in_channels, hidden_dimension, layers, downscaling_factor, num_heads, head_dim, window_size,
                 relative_pos_embedding):
        super().__init__()
        assert layers % 2 == 0, 'Stage layers need to be divisible by 2 for regular and shifted block.'

        self.patch_partition = PatchMerging(in_channels=in_channels, out_channels=hidden_dimension,
                                            downscaling_factor=downscaling_factor)

        self.layers = nn.ModuleList([])
        for _ in range(layers // 2):
            self.layers.append(nn.ModuleList([
                SwinBlock(dim=hidden_dimension, heads=num_heads, head_dim=head_dim, mlp_dim=hidden_dimension * 4,
                          shifted=False, window_size=window_size, relative_pos_embedding=relative_pos_embedding),
                SwinBlock(dim=hidden_dimension, heads=num_heads, head_dim=head_dim, mlp_dim=hidden_dimension * 4,
                          shifted=True, window_size=window_size, relative_pos_embedding=relative_pos_embedding),
            ]))

    def forward(self, x):
        x = self.patch_partition(x)
        for regular_block, shifted_block in self.layers:
            x = regular_block(x)
            x = shifted_block(x)
        return x.permute(0, 3, 1, 2)

class SwinTransformer(nn.Module):
    def __init__(self, *, hidden_dim, layers, heads, channels=3, num_classes=1000, head_dim=32, window_size=7,
                 downscaling_factors=(4, 2, 2, 2), relative_pos_embedding=True):
        super().__init__()

        self.stage1 = StageModule(in_channels=channels, hidden_dimension=hidden_dim, layers=layers[0],
                                  downscaling_factor=downscaling_factors[0], num_heads=heads[0], head_dim=head_dim,
                                  window_size=window_size, relative_pos_embedding=relative_pos_embedding)
        self.stage2 = StageModule(in_channels=hidden_dim, hidden_dimension=hidden_dim * 2, layers=layers[1],
                                  downscaling_factor=downscaling_factors[1], num_heads=heads[1], head_dim=head_dim,
                                  window_size=window_size, relative_pos_embedding=relative_pos_embedding)
        self.stage3 = StageModule(in_channels=hidden_dim * 2, hidden_dimension=hidden_dim * 4, layers=layers[2],
                                  downscaling_factor=downscaling_factors[2], num_heads=heads[2], head_dim=head_dim,
                                  window_size=window_size, relative_pos_embedding=relative_pos_embedding)
        self.stage4 = StageModule(in_channels=hidden_dim * 4, hidden_dimension=hidden_dim * 8, layers=layers[3],
                                  downscaling_factor=downscaling_factors[3], num_heads=heads[3], head_dim=head_dim,
                                  window_size=window_size, relative_pos_embedding=relative_pos_embedding)

        self.mlp_head = nn.Sequential(
            nn.LayerNorm(hidden_dim * 8),
            nn.Linear(hidden_dim * 8, num_classes)
        )

    def forward(self, img):
        x = self.stage1(img)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = x.mean(dim=[2, 3])
        return self.mlp_head(x)

def swin_t(num_classes=100, hidden_dim=96, layers=(2, 2, 6, 2), heads=(3, 6, 12, 24), **kwargs):
    return SwinTransformer(hidden_dim=hidden_dim, layers=layers, heads=heads, **kwargs)

def swin_s(num_classes=100, hidden_dim=96, layers=(2, 2, 18, 2), heads=(3, 6, 12, 24), **kwargs):
    return SwinTransformer(hidden_dim=hidden_dim, layers=layers, heads=heads, **kwargs)

def swin_b(num_classes=100, hidden_dim=128, layers=(2, 2, 18, 2), heads=(4, 8, 16, 32), **kwargs):
    return SwinTransformer(hidden_dim=hidden_dim, layers=layers, heads=heads, **kwargs)

def swin_l(num_classes=100, hidden_dim=192, layers=(2, 2, 18, 2), heads=(6, 12, 24, 48), **kwargs):
    return SwinTransformer(hidden_dim=hidden_dim, layers=layers, heads=heads, **kwargs)


from google.cloud import bigquery
import os
client=bigquery.Client()
os.environ['GOOGLE_APPLICATION_CREDENTIALS']='D:\Pycharm\GCP\gkey-bigquery-practice-sample-31e14021aadb.json'
rows_table=client.list_rows('bigquery-practice-sample.practice_data_manipulation.json_load_tbl')
rows=list(rows_table)
print("the Downloaded {} rows from the table {} ".format(rows,rows_table))
table_id='bigquery-practice-sample.gcp_practice_sample.load_tbl_from_code'
job_config=bigquery.LoadJobConfig(
    schema=[
        bigquery.SchemaField("name","STRING"),
        bigquery.SchemaField("post_abbr","STRING"),
],
    source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON,
)
uri="gs://cloud-samples-data/bigquery/us-states/us-states.json"
load_job=client.load_table_from_uri(
    uri,
    table_id,
    location="US",
    job_config=job_config,
)
load_job.result()
destination_table=client.get_table(table_id)
print("Landed {} rows".format(destination_table.num_rows))


'''
Script for a three-layer neural field simulator.

The scripts sets up an architecture with three one-dimensional neural
fields: two excitatory fields, u and w, and a shared inhibitory field, v.
Fields u and w feature local self-excitation and project to each other
and to field v in a local excitatory fashion. Field v inhibits both
excitatory fields, either locally or globally.

The script then runs the simulator, first with timing, then with online
plotting.
'''

import numpy as np
import matplotlib.pyplot as plt
import time

### choose here whether cosivina is used with or without numba ###

from cosivina.nonumba import *
# from cosivina.numba import *

if options.useNumba:
    print('Using cosvina with numba (just-in-time compilation). Each function '
            'will be compiled when first called, which can take a long time, '
            'but subsequent function calls will be substantially faster.\n')
else:
    print('Using cosvina without numba.\n')

print('Note: Changing between numba and no-numba mode may require restarting the kernel.\n')

## create sim

# Note: We could create the simulator by just loading from the json file:
# sim = Simulator(file = 'presetThreeLayerField_changeDetection.json')

print('Creating simulator object...')

# shared parameters
fieldSize = (1, 180) # important: size for elements must be defined as a tuple
sigma_exc = 5
sigma_inh = 10

# create simulator object
sim = Simulator()

# create inputs (and sum for visualization)
sim.addElement(GaussStimulus1D('stimulus 1', fieldSize, sigma_exc, 0, round(1/4*fieldSize[1]), True, False))
sim.addElement(GaussStimulus1D('stimulus 2', fieldSize, sigma_exc, 0, round(1/2*fieldSize[1]), True, False))
sim.addElement(GaussStimulus1D('stimulus 3', fieldSize, sigma_exc, 0, round(3/4*fieldSize[1]), True, False))
sim.addElement(SumInputs('stimulus sum', fieldSize), ['stimulus 1', 'stimulus 2', 'stimulus 3'])
sim.addElement(ScaleInput('stimulus scale w', fieldSize, 0), 'stimulus sum')

# create neural field
sim.addElement(NeuralField('field u', fieldSize, 20, -5, 4), 'stimulus sum')
sim.addElement(NeuralField('field v', fieldSize, 5, -5, 4))
sim.addElement(NeuralField('field w', fieldSize, 20, -5, 4), 'stimulus scale w')

# # shifted input sum (for plot) - currently not supported in numba mode (h is not a component of NeuralField)
# sim.addElement(SumInputs('shifted stimulus sum', fieldSize), ['stimulus sum', 'field u'], ['output', 'h'])
# sim.addElement(SumInputs('shifted stimulus sum w', fieldSize), ['stimulus scale w', 'field w'], ['output', 'h'])

# create interactions
sim.addElement(GaussKernel1D('u -> u', fieldSize, sigma_exc, 0, True, True), 'field u', 'output', 'field u')
sim.addElement(GaussKernel1D('u -> v', fieldSize, sigma_exc, 0, True, True), 'field u', 'output', 'field v')
sim.addElement(GaussKernel1D('u -> w', fieldSize, sigma_exc, 0, True, True), 'field u', 'output', 'field w')

sim.addElement(GaussKernel1D('v -> u (local)', fieldSize, sigma_inh, 0, True, True), 'field v', 'output', 'field u')
sim.addElement(GaussKernel1D('v -> w (local)', fieldSize, sigma_inh, 0, True, True), 'field v', 'output', 'field w')

# the dimension(s) to sum over must be specified as a numpy array
sim.addElement(SumDimension('sum v', np.array([2]), (1, 1), 1), 'field v', 'output')
sim.addElement(ScaleInput('v -> u (global)', fieldSize, 0), 'sum v', 'output', 'field u')
sim.addElement(ScaleInput('v -> w (global)', fieldSize, 0), 'sum v', 'output', 'field w')

sim.addElement(GaussKernel1D('w -> u', fieldSize, sigma_exc, 0, True, True), 'field w', 'output', 'field u')
sim.addElement(GaussKernel1D('w -> v', fieldSize, sigma_exc, 0, True, True), 'field w', 'output', 'field v')
sim.addElement(GaussKernel1D('w -> w', fieldSize, sigma_exc, 0, True, True), 'field w', 'output', 'field w')

# create noise stimulus and noise kernel
sim.addElement(NormalNoise('noise u', fieldSize, 1));
sim.addElement(GaussKernel1D('noise kernel u', fieldSize, 0, 2, True, True), 'noise u', 'output', 'field u')
sim.addElement(NormalNoise('noise v', fieldSize, 1))
sim.addElement(GaussKernel1D('noise kernel v', fieldSize, 0, 1, True, True), 'noise v', 'output', 'field v')
sim.addElement(NormalNoise('noise w', fieldSize, 1))
sim.addElement(GaussKernel1D('noise kernel w', fieldSize, 0, 2, True, True), 'noise w', 'output', 'field w')

# load parameter settings
sim.loadSettings('presetThreeLayerField_changeDetection.json')

## running the simulation

tMax = 1000

print('Running simulation (first time) ...')
t0 = time.time()
sim.run(tMax, True)
t1 = time.time()
print(f'Time taken: {round(t1-t0, ndigits=3)} seconds')

print('Running simulation (second time) ...')
t0 = time.time()
sim.run(tMax, True)
t1 = time.time()
print(f'Time taken: {round(t1-t0, ndigits=3)} seconds')

## online plotting of field activations

print('\nRunning simulation with online plotting (plotting may be slow in some environments) ...')

# turn on interactive plotting
plt.ion()

# re-initialize
sim.init()

# prepare axes
fig, axes = plt.subplots(3, 1)
for i in range(3):
    axes[i].set_xlim(0, fieldSize[1])
    axes[i].set_ylim(-15, 15)
    axes[i].set_ylabel('activation')
axes[2].set_xlabel('feature space')

# plot initial state
x = np.arange(fieldSize[1])
plot_stim, plot_u = axes[0].plot(x, sim.getComponent('stimulus sum', 'output')[0],
        x, sim.getComponent('field u', 'activation')[0], color='b')
plot_stim.set_color('g')
plot_v, = axes[1].plot(x, sim.getComponent('field v', 'activation')[0], color='b')
plot_stimw, plot_w = axes[2].plot(x, sim.getComponent('stimulus scale w', 'output')[0],
        x, sim.getComponent('field w', 'activation')[0], color='b')
plot_stimw.set_color('g')

# run simulation
tMax = 100
for t in range(tMax):
    sim.step()

    plot_u.set_ydata(sim.getComponent('field u', 'activation')[0])
    plot_v.set_ydata(sim.getComponent('field v', 'activation')[0])
    plot_w.set_ydata(sim.getComponent('field w', 'activation')[0])
    plot_stim.set_ydata(sim.getComponent('stimulus sum', 'output')[0])
    plot_stimw.set_ydata(sim.getComponent('stimulus scale w', 'output')[0])
    fig.canvas.draw()
    fig.canvas.flush_events()
    time.sleep(0.01)


# -*- coding: utf-8 -*-
"""
Created on Tue Aug  4 18:43:42 2020

@author: DELL
"""

import tkinter as tk
import tkinter.font as TkFont
from datetime import datetime

def run():
  current_time=datetime.now()
  diff=current_time - start_time
  txt_var.set('%d.%02d'%(diff.seconds,diff.microseconds//10000))

  if running:
    root.after(20,run)

def start():
  global running
  global start_time

  if not running:
    running=True
    start_time=datetime.now()
    root.after(10,run)

def stop():
  global running
  running=False

def reset():
  global start_time
  start_time=datetime.now()

  if not running:
    txt_var.set('0.00')

running=False
start_time=None

root=tk.Tk()
root.geometry("500x200")
root.title("StopWatch")

txt_var=tk.StringVar()
txt_var.set('0.00')

fontStyle=TkFont.Font(family="Lucida Grade",size=50)
tk.Label(root,textvariable=txt_var,font=fontStyle).pack()

tk.Button(root,text="Start",command=start).pack(fill='x')
tk.Button(root,text="Stop",command=stop).pack(fill='x')
tk.Button(root,text="Reset",command=reset).pack(fill='x')
tk.Button(root,text="Time Lapsed",command=run).pack(fill='x')

root.mainloop()


from torchvision import datasets
from torch.utils.data import DataLoader
import torch.optim as optim
import torch.nn as nn
import torch
from model import lenet

def train(model, criterion, optimizer, train_loader, validation_loader, epochs):
    i=100
    stuff = {'training_loss': [], 'validation_loss': []}

    for epoch in range(epochs):
        for i, (x,y) in enumerate(train_loader):
            optimizer.zero_grad()
            z = model(x.view(-1, 28*28))
            loss = criterion(z, y)
            loss.backward()
            optimizer.step()
            stuff['training_loss'].append(loss.data.item())

        correct = 0
        for (x,y) in validation_loader:
            z = model(x.view(-1, 28*28))
            _, label = torch.max(z)
            correct += (label == y).sum().item()

        accuracy = 100 * (correct / len(validation_dataset))
        stuff['validation_accuracy'].append(accuracy)
    return stuff

if __name__ == '__main__':
    model = lenet.LeNet()

    criterion = nn.CrossEntropyLoss()

    l_rate = 0.01
    optimizer = optim.Adam(model.parameters(), lr=l_rate)

    train_dataset = datasets.MNIST(
        root='datasets/mnist', train=True, download=False)
    validation_dataset = datasets.MNIST(
        root='datasets/mnist', train=False, download=False)

    train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)

    validation_loader = DataLoader(dataset=validation_dataset, batch_size=64, shuffle=True)

    dataset_sizes = {
        'train':len(train_dataset), 'val':len(validation_dataset)
        }

    print(dataset_sizes)
    for batch_ndx, sample in enumerate(train_loader):
        print(sample.inp.is_pinned())
        print(sample.tgt.is_pinned())

    # train(
    #     model=model,
    #     criterion=criterion,
    #     optimizer=optimizer,
    #     train_loader=train_loader,
    #     validation_loader=validation_loader,
    #     epochs=10
    # )


import base64
import json
import logging

from utilities import ONDB_ABORT_IF_UNSUCCESSFUL
from utilities import ONDB_ABSOLUTE
from utilities import ONDB_AP_ALL
from utilities import ONDB_AP_NONE
from utilities import ONDB_AP_SIMPLE_MAJORITY
from utilities import ONDB_CONSISTENCY
from utilities import ONDB_DAYS
from utilities import ONDB_DELETE
from utilities import ONDB_DELETE_IF_VERSION
from utilities import ONDB_DIRECTION
from utilities import ONDB_DURABILITY
from utilities import ONDB_END_INCLUSIVE
from utilities import ONDB_END_VALUE
from utilities import ONDB_FIELD
from utilities import ONDB_FIELD_RANGE
from utilities import ONDB_FORWARD
from utilities import ONDB_HOURS
from utilities import ONDB_INCLUDED_TABLES
from utilities import ONDB_MASTER_SYNC
from utilities import ONDB_MAX_RESULTS
from utilities import ONDB_NONE_REQUIRED
from utilities import ONDB_NONE_REQUIRED_NO_MASTER
from utilities import ONDB_OPERATION
from utilities import ONDB_OPERATION_TYPE
from utilities import ONDB_PERMISSIBLE_LAG
from utilities import ONDB_PUT
from utilities import ONDB_PUT_IF_ABSENT
from utilities import ONDB_PUT_IF_PRESENT
from utilities import ONDB_PUT_IF_VERSION
from utilities import ONDB_RC_ALL
from utilities import ONDB_RC_NONE
from utilities import ONDB_RC_VALUE
from utilities import ONDB_RC_VERSION
from utilities import ONDB_READ_OPTIONS
from utilities import ONDB_REPLICA_ACK
from utilities import ONDB_REPLICA_SYNC
from utilities import ONDB_RETURN_CHOICE
from utilities import ONDB_REVERSE
from utilities import ONDB_ROW
from utilities import ONDB_SIMPLE_CONSISTENCY
from utilities import ONDB_SP_NO_SYNC
from utilities import ONDB_SP_SYNC
from utilities import ONDB_SP_WRITE_NO_SYNC
from utilities import ONDB_START_INCLUSIVE
from utilities import ONDB_START_VALUE
from utilities import ONDB_TABLE_NAME
from utilities import ONDB_TIMEOUT
from utilities import ONDB_TIMEUNIT
from utilities import ONDB_TIME_CONSISTENCY
from utilities import ONDB_TTL_TIMEUNIT
from utilities import ONDB_TTL_VALUE
from utilities import ONDB_UNORDERED
from utilities import ONDB_UPDATE_TTL
from utilities import ONDB_VERSION
from utilities import ONDB_VERSION_CONSISTENCY

from oracle.kv.proxy.gen.ttypes import TConsistency
from oracle.kv.proxy.gen.ttypes import TDirection
from oracle.kv.proxy.gen.ttypes import TDurability
from oracle.kv.proxy.gen.ttypes import TFieldRange
from oracle.kv.proxy.gen.ttypes import TOperation
from oracle.kv.proxy.gen.ttypes import TOperationType
from oracle.kv.proxy.gen.ttypes import TReadOptions
from oracle.kv.proxy.gen.ttypes import TReplicaAckPolicy
from oracle.kv.proxy.gen.ttypes import TReturnChoice
from oracle.kv.proxy.gen.ttypes import TRow
from oracle.kv.proxy.gen.ttypes import TSimpleConsistency
from oracle.kv.proxy.gen.ttypes import TSyncPolicy
from oracle.kv.proxy.gen.ttypes import TTimeConsistency
from oracle.kv.proxy.gen.ttypes import TTimeToLive
from oracle.kv.proxy.gen.ttypes import TTimeUnit
from oracle.kv.proxy.gen.ttypes import TVersionConsistency
from oracle.kv.proxy.gen.ttypes import TWriteOptions

logger = logging.getLogger('nosqldb')

class RestrictedDict(dict):
    """
    dict wrapper that restrict you to add only a predefined set
    of keys. This is for custom parameters that expect certain
    keys.
    """
    def __init__(self, allowed_keys, seq=(), **kwargs):
        """
        This is exactly the same as a regular dict init
        but it has a allowed_keys parameter that let you
        restrict the valid keys in this dict.
        """
        super(RestrictedDict, self).__init__()
        self._allowed_keys = tuple(allowed_keys)
        # normalize arguments to a (key, value) iterable
        if hasattr(seq, 'keys'):
            get = seq.__getitem__
            seq = ((k, get(k)) for k in seq.keys())
        if kwargs:
            from itertools import chain
            seq = chain(seq, kwargs.iteritems())
        # scan the items keeping track of the keys' order
        for k, v in seq:
            self.__setitem__(k, v)

    def __setitem__(self, key, value):
        """Checks if key is a valid one before setting it"""
        if key in self._allowed_keys:
            super(RestrictedDict, self).__setitem__(key, value)
        else:
            raise KeyError("%s is not allowed as key" % key)

    def update(self, e=None, **kwargs):
        """
        Same as regular dict update but using RestrictedDict.__setitem__()
        instead of regular dict.__setitem__()
        """
        try:
            for k in e:
                self.__setitem__(k, e[k])
        except AttributeError:
            for (k, v) in e:
                self.__setitem__(k, v)
        for k in kwargs:
            self.__setitem__(k, kwargs[k])

    def __eq__(self, other):
        """
        Add the comparison for allowed_keys in addition to the regular
        __eq__()
        """
        if other is None:
            return False
        try:
            allowedcmp = (self._allowed_keys == other._allowed_keys)
            if allowedcmp:
                return super(RestrictedDict, self).__eq__(other)
            else:
                return False
        except AttributeError:
            # other is not a RestrictedDict
            return False

    def __ne__(self, other):
        """Equivalent to not __eq__() """
        return not self.__eq__(other)

    def print_allowed_keys(self):
        """Print the set of allowed keys"""
        return 'Allowed Keys(%s)' % (self._allowed_keys.__repr__())

    def validate_return_opts(self, ret_opts):
        if (ret_opts is not None):
            if((type(ret_opts) is str or type(ret_opts) is unicode) and
                (ret_opts == ONDB_RC_ALL or
                ret_opts == ONDB_RC_NONE or
                ret_opts == ONDB_RC_VALUE or
                ret_opts == ONDB_RC_VERSION)):
                    return True
            else:
                return False
        else:
            return True

    def validate(self):
        """This method should be filled in children"""
        pass

class DataManager:
    # This class is in charge of making the appropiate changes from JSON to
    # the native Thrift classes needed for the protocol.
    @staticmethod
    def dict_to_trow(in_dict):
        """
          Convert from a native python dict to a TRow.
          @param in_dict: the dictoinary to be converted to TRow.
          @return: a TRow with the desired information.
        """
        t_ttl = None
        if (in_dict is not None):
            res = json.dumps(in_dict)
            try:
                ttl = in_dict.get_timetolive()
                if (ttl is not None):
                    val = ttl.get(ONDB_TTL_VALUE)
                    timeunit = ttl.get(ONDB_TTL_TIMEUNIT)
                    t_timeunit = DataManager.from_json_to_ttimeunit(timeunit)
                    t_ttl = TTimeToLive(value=val, timeUnit=t_timeunit)
            except AttributeError:
                pass
        else:
            res = {}
        return TRow(jsonRow=res, ttl=t_ttl)

    @staticmethod
    def encode_binary(data):
        # Encode to Base64 the given data
        if (data is not None):
            return base64.b64encode(data)
        else:
            return None

    @staticmethod
    def decode_binary(data):
        # Decode from Base64 the given data
        if (data is not None):
            return base64.b64decode(data)
        else:
            return None

    @staticmethod
    def trow_to_dict(trow):
        # Convert back from a TRow to a dict with native python values.
        # @param twriteresult: the resulting data from the put operation.
        # @return: a dictionary with the data returned.
        if (trow is not None):
            res = json.loads(trow.jsonRow)
        else:
            res = None
        return res

    @staticmethod
    def from_json_to_tsync_policy(policy):
        # This method converts from synch policy to the Thrift equivalent
        if (policy is not None):
            if (policy == ONDB_SP_NO_SYNC):
                return TSyncPolicy.NO_SYNC
            elif (policy == ONDB_SP_SYNC):
                return TSyncPolicy.SYNC
            elif (policy == ONDB_SP_WRITE_NO_SYNC):
                return TSyncPolicy.WRITE_NO_SYNC
        return None

    @staticmethod
    def from_json_to_treplica_ack_policy(policy):
        # This method converts from replica ack policy to Thrift equivalent
        if (policy is not None):
            if (policy == ONDB_AP_ALL):
                return TReplicaAckPolicy.ALL
            elif (policy == ONDB_AP_NONE):
                return TReplicaAckPolicy.NONE
            elif (policy == ONDB_AP_SIMPLE_MAJORITY):
                return TReplicaAckPolicy.SIMPLE_MAJORITY
        return None

    @staticmethod
    def from_json_to_tdurability(durability):
        # This method converts from Durability to the Thrift equivalent
        master_sync_val = replica_sync_val = replica_ack_val = None
        if (durability is not None and isinstance(durability, dict)):
            master_sync_val = DataManager.from_json_to_tsync_policy(
                durability.get(ONDB_MASTER_SYNC, None))
            replica_sync_val = DataManager.from_json_to_tsync_policy(
                durability.get(ONDB_REPLICA_SYNC, None))
            replica_ack_val = DataManager.from_json_to_treplica_ack_policy(
                durability.get(ONDB_REPLICA_ACK, None))
        logger.debug("master_sync={0} replica_sync={1} replica_ack={2}".
            format(master_sync_val, replica_sync_val, replica_ack_val))
        # return something as long as there is relevant information to encode
        if (master_sync_val is None and
           replica_sync_val is None and
           replica_ack_val is None):
            return None
        else:
            return TDurability(masterSync=master_sync_val,
                               replicaSync=replica_sync_val,
                               replicaAck=replica_ack_val)

    @staticmethod
    def from_json_to_treturn_choice(return_opt):
        # This method converts from Return Choice to the Thrift equivalent
        if (return_opt is not None):
            if (return_opt == ONDB_RC_NONE):
                return TReturnChoice.NONE
            elif (return_opt == ONDB_RC_ALL):
                return TReturnChoice.ALL
            elif (return_opt == ONDB_RC_VALUE):
                return TReturnChoice.ONLY_VALUE
            elif (return_opt == ONDB_RC_VERSION):
                return TReturnChoice.ONLY_VERSION
        return None

    @staticmethod
    def from_json_to_twrite_options(write_opts):
        # This method converts from Write Options to the Thrift equivalent
        durability_val = timeout_val = return_opt_val = None
        if (write_opts is not None):
            # encode durability
            dur = write_opts.get(ONDB_DURABILITY, None)
            if (dur is not None):
                durability_val = DataManager.from_json_to_tdurability(dur)
            else:
                durability_val = None
            # encode timeout
            timeout_val = int(write_opts.get(ONDB_TIMEOUT, 0))
            # encode the Return Choice
            return_opt_val = DataManager.from_json_to_treturn_choice(
                write_opts.get(ONDB_RETURN_CHOICE, None))
            update_val = write_opts.get(ONDB_UPDATE_TTL)
        # return something as long as there is something relevant to encode
        if (durability_val is None and timeout_val is None and
            return_opt_val is None and update_val is None):
            return None
        else:
            return TWriteOptions(durability=durability_val,
                                 timeoutMs=timeout_val,
                                 returnChoice=return_opt_val,
                                 updateTTL=update_val)

    @staticmethod
    def from_json_to_tsimple_consistency(consistency):
        # This method converts from simple consistency to the Thrift equivalent
        if (consistency is not None):
            if (consistency == ONDB_NONE_REQUIRED):
                return TConsistency(simple=TSimpleConsistency.NONE_REQUIRED)
            elif (consistency == ONDB_ABSOLUTE):
                return TConsistency(simple=TSimpleConsistency.ABSOLUTE)
            elif (consistency == ONDB_NONE_REQUIRED_NO_MASTER):
                return TConsistency(
                    simple=TSimpleConsistency.NONE_REQUIRED_NO_MASTER)
        return None

    @staticmethod
    def from_json_to_ttime_consistency(t_consistency):
        # This method converts from TimeConsistency to Thrift equivalent
        if (t_consistency is not None):
            lag = t_consistency.get(ONDB_PERMISSIBLE_LAG, None)
            to = t_consistency.get(ONDB_TIMEOUT, None)
            return TConsistency(
                time=TTimeConsistency(permissibleLag=lag, timeoutMs=to))
        return None

    @staticmethod
    def from_json_to_tversion_consistency(v_consistency):
        # This method converts from VersionConsistency to Thrift equivalent
        if (v_consistency is not None):
            ver = v_consistency.get(ONDB_VERSION, None)
            to = v_consistency.get(ONDB_TIMEOUT, None)
            return TConsistency(
                version=TVersionConsistency(version=ver, timeoutMs=to))
        return None

    @staticmethod
    def from_json_to_tread_options(read_opts):
        # This method converts from Read Options to the Thrift equivalent
        consistency_val = timeout_val = None
        if (read_opts is not None and type(read_opts) is dict):
            # encode Consistency
            con = read_opts.get(ONDB_CONSISTENCY, None)
            if (con is not None):
                simple_con = con.get(ONDB_SIMPLE_CONSISTENCY, None)
                if (simple_con is not None):
                    consistency_val = \
                        DataManager.from_json_to_tsimple_consistency(simple_con)
                else:
                    time_con = read_opts.get(ONDB_TIME_CONSISTENCY, None)
                    if (time_con is not None):
                        consistency_val = \
                            DataManager.from_json_to_ttime_consistency(time_con)
                    else:
                        version_con = read_opts.get(
                            ONDB_VERSION_CONSISTENCY, None)
                        if (version_con is not None):
                            consistency_val = \
                                DataManager.from_json_to_tversion_consistency(
                                    version_con)
            # encode timeout
            timeout_val = int(read_opts.get(ONDB_TIMEOUT, 0))
        # return something as long as there is something to encode
        if (consistency_val is None and timeout_val is None):
            return None
        else:
            return TReadOptions(consistency=consistency_val,
                                timeoutMs=timeout_val)

    @staticmethod
    def from_json_to_tfield_range(field_range):
        # This method converts from Field Range to the Thrift equivalent
        if (field_range is not None):
            # get all the relevant information
            field = field_range.get(ONDB_FIELD, None)
            start_value = field_range.get(ONDB_START_VALUE, None)
            end_value = field_range.get(ONDB_END_VALUE, None)
            start_inclusive = field_range.get(ONDB_START_INCLUSIVE, None)
            end_inclusive = field_range.get(ONDB_END_INCLUSIVE, None)
            # encode it as a TFieldRange
            return TFieldRange(fieldName=field,
                               startValue=start_value,
                               startIsInclusive=start_inclusive,
                               endValue=end_value,
                               endIsInclusive=end_inclusive)
        return None

    @staticmethod
    def from_json_to_tmultirow_options(multirow_opts):
        # This method converts from Multirow Options to the Thrift equivalent
        t_field_range = None
        t_tables = None
        if (multirow_opts is not None):
            # get field range
            t_field_range = DataManager.from_json_to_tfield_range(
                multirow_opts.get(ONDB_FIELD_RANGE, None))
            # get the list of parent and children tables and consolidate
            t_tables = multirow_opts.get(ONDB_INCLUDED_TABLES, None)
        # return the data
        return t_field_range, t_tables

    @staticmethod
    def from_json_to_tdirection(direction):
        # This method converts from Direction to the Thrift equivalent
        res = None
        if (isinstance(direction, dict)):
            d = direction.get(ONDB_DIRECTION, None)
            if (d == ONDB_FORWARD):
                res = TDirection.FORWARD
            elif (d == ONDB_REVERSE):
                res = TDirection.REVERSE
            elif (d == ONDB_UNORDERED):
                res = TDirection.UNORDERED
        return res

    @staticmethod
    def from_json_to_ttable_iterator_options(table_iterator_opts):
        # This method converts from Table Iterator Options to the Thrift
        # equivalent
        t_dir = None
        max_results = 0
        tr_opts = None
        if (table_iterator_opts is not None):
            # get direction, max results and read options
            direction = table_iterator_opts.get(ONDB_DIRECTION, None)
            max_results_by_batch = table_iterator_opts.get(
                ONDB_MAX_RESULTS,
                0)
            read_opts = table_iterator_opts.get(ONDB_READ_OPTIONS, None)
            t_dir = DataManager.from_json_to_tdirection(direction)
            if (max_results_by_batch == 0 or max_results_by_batch is None):
                max_results = 0
            else:
                max_results = int(max_results_by_batch)
            tr_opts = DataManager.from_json_to_tread_options(read_opts)
        # return the Thrift data
        return t_dir, max_results, tr_opts

    @staticmethod
    def from_json_to_toperation_type(operation_type):
        # This methid converts from Operation Type
        op_type = operation_type.get(ONDB_OPERATION_TYPE, None)
        if (op_type == ONDB_DELETE):
            return TOperationType.DELETE
        elif (op_type == ONDB_DELETE_IF_VERSION):
            return TOperationType.DELETE_IF_VERSION
        elif (op_type == ONDB_PUT):
            return TOperationType.PUT
        elif (op_type == ONDB_PUT_IF_ABSENT):
            return TOperationType.PUT_IF_ABSENT
        elif (op_type == ONDB_PUT_IF_PRESENT):
            return TOperationType.PUT_IF_PRESENT
        elif (op_type == ONDB_PUT_IF_VERSION):
            return TOperationType.PUT_IF_VERSION
        return None

    @staticmethod
    def from_json_to_toperations(table_operations):
        # This method converts from Operations list to the Thrift equivalent
        operations = []
        if (isinstance(table_operations, list)):
            for oper in table_operations:
                # get table name, operation type, abort option, row,
                # return choice and version
                tn = oper.get(ONDB_TABLE_NAME, None)
                op_str = oper.get(ONDB_OPERATION, None)
                op = DataManager.from_json_to_toperation_type(
                    op_str)
                abort = oper.get(ONDB_ABORT_IF_UNSUCCESSFUL, False)
                rd = oper.get(ONDB_ROW, None)
                row = DataManager.dict_to_trow(rd)
                r_choice = DataManager.from_json_to_treturn_choice(
                    rd.get(ONDB_RETURN_CHOICE, None))
                ver = oper.get(ONDB_VERSION, None)
                # convert it as TOperation and append it to the list
                t_oper = TOperation(tn, op, row, r_choice, abort, ver)
                operations.append(t_oper)
        # return the list
        return operations

    @staticmethod
    def from_json_to_ttimeunit(timeunit):
        # This method converts from TimeUnit to the Thrift equivalent
        tu = timeunit.get(ONDB_TIMEUNIT, None)
        if (tu == ONDB_HOURS):
            return TTimeUnit.HOURS
        elif (tu == ONDB_DAYS):
            return TTimeUnit.DAYS
        return None


# -*- coding: utf-8 -*-
"""
Created on Tue Dec  1 19:04:45 2020

@author: Linter
"""
from art import logo
print(logo)
action_dict = {}
while True:
    name = input("what's your name: " )
    bid = float(input("what's your bid: " ))
    ask_more = input("Are there any other bidders? Type 'yes' or 'no'. \n")
    action_dict[name] = bid
    print(action_dict)

    if ask_more == "yes":
        continue
    elif ask_more == "no":
        max_bid = 0
        for key in action_dict:
            if max_bid < action_dict[key]:
                max_bid = action_dict[key]
        winner_index = list(action_dict.values()).index(max_bid)
        winner = list(action_dict.keys())[winner_index]
        print(f"The winner is {winner} ,bid is {max_bid}")
        break


import numpy as onp
from torch.utils import data

def numpy_collate(batch):
    if isinstance(batch[0], onp.ndarray):
        return onp.stack(batch)
    elif isinstance(batch[0], (tuple, list)):
        transposed = zip(*batch)
        return [numpy_collate(samples) for samples in transposed]
    else:
        return onp.array(batch)

class NumpyLoader(data.DataLoader):
    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0,
                 pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None):
        super(self.__class__, self).__init__(dataset,
                                             batch_size=batch_size,
                                             shuffle=shuffle,
                                             sampler=sampler,
                                             batch_sampler=batch_sampler,
                                             num_workers=num_workers,
                                             collate_fn=numpy_collate,
                                             pin_memory=pin_memory,
                                             drop_last=drop_last,
                                             timeout=timeout,
                                             worker_init_fn=worker_init_fn
                                             )


import asyncio

async def handle_client(reader, writer):
    data = await reader.read(100)
    message = data.decode("utf-8")
    addr = writer.get_extra_info("peername")

    print(f"Received {message} from {addr}")

    print("Send: ACK")
    writer.write(b"ACK")
    await writer.drain()

    print("Closing the connection")
    writer.close()
    await writer.wait_closed()

async def main():
    server = await asyncio.start_server(handle_client, "127.0.0.1", 8888)

    addr = server.sockets[0].getsockname()
    print(f"Serving on {addr}")

    async with server:
        await server.serve_forever()

asyncio.run(main())