#!/usr/bin/env python3
from guizero import App, Combo, PushButton, Text
from screeninfo import get_monitors
from python_mpv_jsonipc import MPV
from threading import Thread
import socket
import subprocess
import os
import json
import time
import numpy as np
os.environ["DISPLAY"] = ":0.0"

sdobSocketKiller = 1
def thread_sdobSocket():
    global sdobSocketKiller
    global csock
    global dataSet
    print('Start Thread Socket')
    while sdobSocketKiller:
        try:
            (bytes, address) = csock.recvfrom(8196)
            msg = bytes.decode('utf-8')
            print('address:', address, 'recv', msg)
            # msgJson = json.loads(msg)
            # if type(msgJson) is dict:
            #    for k in msgJson.keys():
            #        setattr(dataSet, k, msgJson[k])
            #        print(k, msgJson[k])
            #else:
            #    setattr(dataSet, "socketdata", msg)
        except:
            time.sleep(0.2)
    # print('Stop Thread Socket')

def thread_getFrames(file):
    global dataSet
    global statusText
    statusText.value = "Getting KeyFrame Locations"
    cmd = "ffprobe -loglevel error -select_streams v:0 -show_entries packet=pts_time,flags -of csv=print_section=0 \"{}\" | awk -F',' '/K/ {{print $1}}'".format(file)
    proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)
    frameList = []
    for line in proc.stdout.readlines():
        frameList.append(line.decode('utf-8').rstrip())
    setattr(dataSet, "frameList", frameList)
    statusText.value = "Keyframes Found, Finished Looking"

def thread_processVideo():
    global dataSet
    global statusText

    videoRoot = '/tmp'
    videoFiles = []
    videoExt = '.mp4'
    videoDest = '/home/pi/Videos/120.mp4'

    if (hasattr(dataSet, "fileext")):
        videoExt = getattr(dataSet, "fileext")

    if (hasattr(dataSet, "slate")):
        # Slate
        sTime = float(getattr(dataSet, "slate")) - 2.0
        if (sTime < 0):
            sTime = 0.00
        eTime = float(getattr(dataSet, "slate")) + 5.0
        sTimes = splitVideoFileKeyTimes(sTime, eTime)
        # print("Slate Start Adding", sTime - sTimes[0])
        # print("Slate End Adding", sTimes[1] - eTime)
        videoFileName = 'sdobSlate{}'.format(videoExt)
        videoFile = '{}/{}'.format(videoRoot, videoFileName)
        videoFiles.append(videoFileName)
        splitVideoFile(getattr(dataSet, "filename"), sTimes[0], sTimes[1], videoFile)

    if (hasattr(dataSet, "exit")):
        statusText.value = "Splitting Video Skydive"
        # Skydive
        sTime = float(getattr(dataSet, "exit")) - 2.0
        if (sTime < 0):
            sTime = 0.00
        eTime = float(getattr(dataSet, "exit")) + 40.0
        sTimes = splitVideoFileKeyTimes(sTime, eTime)
        # print("Skydive Start Adding", sTime - sTimes[0])
        # print("Skydive End Adding", sTimes[1] - eTime)
        videoFileName = 'sdobExit{}'.format(videoExt)
        videoFile = '{}/{}'.format(videoRoot, videoFileName)
        videoFiles.append(videoFileName)
        splitVideoFile(getattr(dataSet, "filename"), sTimes[0], sTimes[1], videoFile)

    statusText.value = "Combining Video Slate and Skydive"
    # Recombine
    combineVideoFiles(videoRoot, videoFiles, videoDest)
    statusText.value = "Done Processing Video"

def globalClicked():
    if (hasattr(dataSet, "playing") and getattr(dataSet, "playing") == "1"):
        mpv_player.stop()

def selectButton(t):
    if (hasattr(dataSet, "playing") and getattr(dataSet, "playing") == "1"):
        return

    if (hasattr(dataSet, "team") and hasattr(dataSet, "round") and getattr(dataSet, "round") == t + 1):
        rX = getattr(r, "round_{}".format(str(t)))
        rX.bg = "blue"
        setattr(r, "round_{}".format(str(t)), rX)
        setattr(dataSet, "round", t + 1)
        print('Playing {} {}'.format(getattr(dataSet, "team"), getattr(dataSet, "round")))
    else:
        rKeys = r.__dict__.keys()
        for k in r.__dict__.keys():
            rX = getattr(r, k)
            rX.bg = "white"
            setattr(r, k, rX)

        rX = getattr(r, "round_{}".format(str(t)))
        rX.bg = "orange"
        setattr(r, "round_{}".format(str(t)), rX)
        setattr(dataSet, "round", t + 1)

def teamChanged(t):
    setattr(dataSet, "team", t)

def submitButton(t, r):
    # print("Submitting", getattr(dataSet, "filename"), getattr(dataSet, "slate"), getattr(dataSet, "exit"))
    if (hasattr(dataSet, "playing") and getattr(dataSet, "playing") == "1"):
        return

    if (hasattr(dataSet, "filename") and hasattr(dataSet, "team") and hasattr(dataSet, "round")):
        print('Uploading For {:s} Round: {:d}'.format(getattr(dataSet, "team"), getattr(dataSet, "round")))

    if (hasattr(dataSet, "filename") and hasattr(dataSet, "slate") and hasattr(dataSet, "exit")):
        print('Splitting Video At', getattr(dataSet, "slate"), 'and', getattr(dataSet, "exit"));

        # Start to Stop
        # splitVideoFile(getattr(dataSet, "filename"), getattr(dataSet, "slate"), getattr(dataSet, "exit"), "/tmp/test1.mp4")
        statusText.value = "Processing Video"
        processVideoThread = Thread(target = thread_processVideo)
        processVideoThread.setDaemon(True)
        processVideoThread.start()

def combineVideoFiles(rootFolder, videoArr, outputFilename):
    videoListPath = rootFolder + "/sdobVideoList.txt"
    f = open(videoListPath, "w")
    for v in videoArr:
        f.write("file '{}'\n".format(v))
    f.close()
    vidCmd = 'ffmpeg -y -f concat -i  "{}" -c copy {}'.format(videoListPath, outputFilename)
    os.system(vidCmd)
    # proc = subprocess.Popen(vidCmd, shell=True, stderr=subprocess.PIPE)
    # for line in proc.stderr.readlines():
    #     statusText.value = line.decode('utf-8').rstrip()

def splitVideoFileKeyTimes(startingTime, endingTime):
    ret = np.zeros(2, dtype = float)
    startTime = 0.00
    endTime = 0.00
    checkStart = float(startingTime)
    checkEnd = float(endingTime)
    # print(startingTime, endingTime)
    for f in getattr(dataSet, "frameList"):
        # print('Keyframe', f)
        if (float(f) <= checkStart):
            ret[0] = float(f)
        if (endTime == 0.00 and float(f) >= checkEnd):
            ret[1] = float(f)
        if (ret[0] != 0.00 and ret[1] != 0.00):
            break
    return ret

def splitVideoFile(filename, startTime, endTime, outputFilename):
    extraArgs = ""
    if (endTime != 0.00):
        extraArgs = extraArgs + " -t " + str(endTime - startTime)

    # print("StartKey", startTime, "EndKey", endTime)
    vid1Cmd = 'ffmpeg -y -i "{}" -ss {} {} -c copy {}'.format(filename, startTime, extraArgs, outputFilename)
    os.system(vid1Cmd)
    statusText.value = "Done!"
    # print(vid1Cmd)
    # proc = subprocess.Popen(vid1Cmd, shell=True, stderr=subprocess.PIPE)
    # for line in proc.stderr.readlines():
    #     statusText.value = line.decode('utf-8').rstrip()

def playnewVideo():
    mpv_player.loop = False
    mpv_player.play("/home/pi/Videos/120.mp4")
    csock.sendto(str.encode('{"event":"playing"}'), ssock_file)

def openButton():
    if (hasattr(dataSet, "playing") and getattr(dataSet, "playing") == "1"):
        return

    filename = app.select_file(folder="/media/", filetypes=[["Media Files", "*.mp4 *.mov *.vob *.wmv *.mpg *.mpeg *.mkv *.m4v *.avi *.ts *.webm"], ["All Files", "*.*"]])
    if (isinstance(filename, str)):
        setattr(dataSet, "filename", filename)
        openvideo_btn.text = os.path.basename(filename)
        openvideo_btn.bg = "#cccc00"
        try:
            fExt = os.path.splitext(filename)[1]
            setattr(dataSet, "fileext", fExt)
        except:
            pass

        keyFrameThread = Thread(target = thread_getFrames, args=[filename])
        keyFrameThread.setDaemon(True)
        keyFrameThread.start()
        try:
            csock.sendto(str.encode('{{"event":"filechange","filename":"{}"}}'.format(filename)), ssock_file)
        except:
            print("Cannot send to sdobox.socket")

    else:
    # if (isinstance(filename, (unicode, tuple))):
        openvideo_btn.text = "Open Video"
        openvideo_btn.bg = "#cc33cc"

def playButton():
    if (hasattr(dataSet, "playing") and getattr(dataSet, "playing") == "1"):
        return

    if (hasattr(dataSet, "filename")):
        mpv_player.loop = False
        mpv_player.play(getattr(dataSet, "filename"))
        try:
            csock.sendto(str.encode('{"event":"playing"}'), ssock_file)
        except:
            print("Cannot send playing to sdobox.socket")

class DataSet(object):
    pass

class TeamBtns(object):
    pass

class RoundBtns(object):
    pass

r = RoundBtns()
dataSet = DataSet()
selectedTeam = ""
selectedRound = 0

mpv_player = MPV(start_mpv=False, ipc_socket="/tmp/mpv.socket")
monitorList = get_monitors()
for m in monitorList:
    print(str(m))

appWidth = 720 # monitorList[0].width
appHeight = 480 # monitorList[0].height
appMulti = appWidth / 640
if (appMulti > 1):
    fontSize = int(26 * (appMulti * .85))
    fontSizeMin = int(16 * (appMulti * .85))
else:
    fontSize = 26
    fontSizeMin = 16

app = App(title="Skydive or Bust Dubbing Station", width=appWidth, height=appHeight, layout="grid")

# disc_combo = Combo(app, options=["Select Discipline", "2way Sequential ProAm", "2way Sequential Open", "4way Sequential Open", "4way Rotations"], command=teamChanged, grid=[0,1,8,1], align="top", width="23")
# disc_combo.text_size = fontSize
# disc_combo.bg = "#cc3333"

# team_combo = Combo(app, options=["Select Team", "Team 1", "Team 2", "Team 3", "Team 4"], command=teamChanged, grid=[0,2,6,1], align="top", width="23")
# team_combo.text_size = fontSize
# team_combo.bg = "#33cc33"

openvideo_btn = PushButton(app, command=openButton, text="Open Video", width="17", grid=[0,4,3,1])
openvideo_btn.text_size = fontSizeMin
openvideo_btn.bg = "#cc33cc"

playvideo_btn = PushButton(app, command=playButton, text="Play Video", width="17", grid=[3,4,3,1])
playvideo_btn.text_size = fontSizeMin

# for x in range(0, 6):
#     rX = PushButton(app, command=selectButton, args=[x], text=str(x + 1), grid=[x, 5])
#     rX.text_size = fontSize
#     rX.width = 3
#     rX.bg = "white"
#     setattr(r, "round_{}".format(str(x)), rX)

# for x in range(6, 10):
#     rX = PushButton(app, command=selectButton, args=[x], text=str(x + 1), grid=[x - 6, 6])
#     rX.text_size = fontSize
#     rX.width = 3
#     rX.bg = "white"
#     setattr(r, "round_{}".format(str(x)), rX)

update_btn = PushButton(app, command=submitButton, args=[selectedTeam, selectedRound], text="Upload Video", grid=[0,7,3,1])
update_btn.text_size = fontSize
update_btn.disabled = True

playnew_btn = PushButton(app, command=playnewVideo, text="Play New Video", grid=[4,7,3,1])
playnew_btn.text_size = fontSize
playnew_btn.disabled = True

# move_btn = PushButton(app, command=submitButton, args=[selectedTeam, selectedRound], text="Move Video", grid=[3,7,3,1])
# move_btn.text_size = fontSize
# move_btn.disabled = True

statusText = Text(app, text="", align="left", size=14, grid=[0,10,7,1])

# Bind to key press events with a decorator
@mpv_player.on_event("start-file")
def evStartFile_handler(evData):
    setattr(dataSet, "playing", "1")

@mpv_player.on_event("end-file")
def evStartFile_handler(evData):
    setattr(dataSet, "playing", "0")

ssock_file = '/tmp/sdobox.socket';
csock_file = '/tmp/sdobox.dubbing.socket';

if os.path.exists(csock_file):
    os.remove(csock_file)

csock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
csock.bind(csock_file)
csock.setblocking(0)
csock.sendto(str.encode('{"event":"started"}'), ssock_file)

sdobThread = Thread(target = thread_sdobSocket)
sdobThread.setDaemon(True)
sdobThread.start()

app.when_clicked = globalClicked
app.display()

sdobSocketKiller = 0;

try:
    mpv_player.terminate()
except:
    pass

try:
    csock.sendto(str.encode('{"event":"stopped"}'), ssock_file)
except:
    pass
finally:
    csock.close()

if os.path.exists(csock_file):
    os.remove(csock_file)

sdobThread.join()


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

class Solution:
    def get_maximum_binary_tree(nums: List[int], start: int, end: int) -> Optional[TreeNode]:
        if end <= start:
            return None
        maximal_i = start
        maximal = 0
        for i in range(start, end):
            if maximal <= nums[i]:
                maximal = nums[i]
                maximal_i = i
        node = TreeNode(maximal)
        node.left = Solution.get_maximum_binary_tree(nums, start, maximal_i)
        node.right = Solution.get_maximum_binary_tree(nums, maximal_i + 1, end)
        return node

    def constructMaximumBinaryTree(self, nums: List[int]) -> Optional[TreeNode]:
        return Solution.get_maximum_binary_tree(nums, 0, len(nums))


class BinarySearch:

    def __init__(self, a):
        self.a = a

    def searchElementPosition(self, h, l, val):
        mid = int((l + h) / 2)
        if self.a[mid] == val:
            return mid
        elif h < l:
            return None
        else:
            if self.a[mid] > val:
                h = mid - 1
            else:
                l = mid + 1
            return self.searchElementPosition(h, l, val)

if __name__ == "__main__":
    listData = [3, 6, 8, 12, 14, 17, 25, 29.31, 36, 42, 47, 53, 55, 62]

    t1 = BinarySearch(listData)
    l = 0
    h = len(listData)
    print(t1.searchElementPosition(h, l, 62))


from typing import List
import pytest
from ledger import TransactionLog, Transaction, AccountNotFoundException

def test_create_account_with_500_balance():
    transaction_log = TransactionLog()
    transaction_log.create_account('john wick', 500.)

    assert transaction_log.get_account_balance('john wick') == 500.

def test_create_account_with_zero_balance():
    transaction_log = TransactionLog()
    transaction_log.create_account('james wick')

    assert transaction_log.get_account_balance('james wick') == 0.

def test_add_transaction():
    transaction_log = TransactionLog()
    transaction_log.create_account('john')
    transaction_log.create_account('james')
    tx = Transaction(tx_date = '2021-5-24 11:00', src_acct = 'john', dst_acct = 'james', tx_value = 100.)
    transaction_log.add_transaction(tx)

    assert transaction_log.get_account_balance('james') == 100
    assert transaction_log.get_account_balance('john') == -100

def test_get_balance_at_date():
    transaction_log = TransactionLog()
    transaction_log.create_account('john')
    transaction_log.create_account('james')

    transaction_log.add_transaction(Transaction(src_acct='john', dst_acct='james', tx_date = '2021-1-10 08:00', tx_value = 100))
    transaction_log.add_transaction(Transaction(src_acct='john', dst_acct='james', tx_date = '2021-2-10 08:00', tx_value = 100))
    transaction_log.add_transaction(Transaction(src_acct='john', dst_acct='james', tx_date = '2021-3-10 08:00', tx_value = 100))

    assert transaction_log.get_account_balance_at('james', '2021-2-25 12:00') == 200.

def test_invalid_account_raises_account_not_found():
    transaction_log = TransactionLog()
    with pytest.raises(AccountNotFoundException):
        transaction_log.get_account_balance('hendrix')

def test_get_transactions_for_account():
    tx1 = Transaction(src_acct='john', dst_acct='james', tx_date = '2021-1-10 08:00', tx_value = 100)
    tx2 = Transaction(src_acct='peter', dst_acct='sarah', tx_date = '2021-2-10 08:00', tx_value = 50)
    transaction_log = TransactionLog([tx1, tx2])

    assert transaction_log.get_transactions('peter') == [tx2]

def test_get_transactions_for_account_to_date():
    tx1 = Transaction(src_acct='john', dst_acct='james', tx_date = '2021-1-1 08:00', tx_value = 50)
    tx2 = Transaction(src_acct='peter', dst_acct='sarah', tx_date = '2021-1-2 08:00', tx_value = 50)
    tx3 = Transaction(src_acct='john', dst_acct='james', tx_date = '2021-1-3 08:00', tx_value = 50)
    tx4 = Transaction(src_acct='peter', dst_acct='sarah', tx_date = '2021-1-4 08:00', tx_value = 50)
    transaction_log = TransactionLog([tx1, tx2, tx3, tx4])

    assert transaction_log.get_transactions_to_date('james', '2021-1-2 23:00') == [tx1]


from itertools import zip_longest
import logging

import numpy as np

from hyperspy.misc.axis_tools import check_axes_calibration
from hyperspy.misc.array_tools import are_aligned

_logger = logging.getLogger(__name__)

def _get_shapes(am, ignore_axis):
    if ignore_axis is not None:
        try:
            ignore_axis = am[ignore_axis]
        except ValueError:
            pass
    sigsh = tuple(axis.size if (ignore_axis is None or axis is not
                                ignore_axis)
                  else 1 for axis in am.signal_axes) if am.signal_dimension != 0 else ()

    navsh = tuple(axis.size if (ignore_axis is None or axis is not
                                ignore_axis)
                  else 1 for axis in am.navigation_axes) if am.navigation_dimension != 0 else ()
    return sigsh, navsh

def are_signals_aligned(*args, ignore_axis=None):
    if len(args) < 2:
        raise ValueError(
            "This function requires at least two signal instances"
        )

    args = list(args)
    am = args.pop().axes_manager

    sigsh, navsh = _get_shapes(am, ignore_axis)
    while args:
        amo = args.pop().axes_manager
        sigsho, navsho = _get_shapes(amo, ignore_axis)

        if not (are_aligned(sigsh[::-1], sigsho[::-1]) and
                are_aligned(navsh[::-1], navsho[::-1])):
            return False
    return True

def _check_and_get_longest_axis(axes):
    """Return the longest axis from a list of axes.

    In the case of ties, the first element in the list
    will be returned. Logs a warning if the axes have
    different calibrations.
    """
    only_left = filter(lambda x: x is not None, axes)
    longest_idx, longest = sorted(
        enumerate(only_left),
        key=lambda x: x[1].size,
        reverse=True,
    )[0]

    # Exit early if not DataAxis objects - None is used
    # in broadcast_signals() below as a filler so we
    # skip it here.
    def _check_wrapper(ax1, ax2):
        if ax1 is None or ax2 is None:
            return True

        return check_axes_calibration(ax1, ax2)

    # Returns a list of bools, where False
    # indicates an axis with a different
    # calibration to the longest axis
    axes_check = [
        _check_wrapper(axes[longest_idx], axes[i])
        for i in range(len(axes))
        if i != longest_idx
    ]

    if not all(axes_check):
        _logger.warning(
            "Axis calibration mismatch detected along axis "
            f"{longest.index_in_axes_manager}. The "
            f"calibration of signal {longest_idx} along "
            "this axis will be applied to all signals "
            "after stacking."
        )

    return longest

def broadcast_signals(*args, ignore_axis=None):
    """Broadcasts signals according to the HyperSpy broadcasting rules.

    signal and navigation spaces are each separately broadcasted according to
    the numpy broadcasting rules. One axis can be ignored and left untouched
    (or set to be size 1) across all signals.

    Parameters
    ----------
    *args : BaseSignal
        Signals to broadcast together
    ignore_axis : {None, str, int, Axis}
        The axis to be ignored when broadcasting

    Returns
    -------
    list of signals

    """
    from hyperspy.signal import BaseSignal

    if len(args) < 2:
        raise ValueError(
            "This function requires at least two signal instances"
            )
    if any([not isinstance(a, BaseSignal) for a in args]):
        raise ValueError("Arguments must be of signal type.")
    args = list(args)
    if not are_signals_aligned(*args, ignore_axis=ignore_axis):
        raise ValueError("The signals cannot be broadcasted")
    else:
        if ignore_axis is not None:
            for s in args:
                try:
                    ignore_axis = s.axes_manager[ignore_axis]
                    break
                except ValueError:
                    pass
        new_nav_axes = []
        new_nav_shapes = []
        for axes in zip_longest(
            *[s.axes_manager.navigation_axes for s in args], fillvalue=None
        ):
            longest = _check_and_get_longest_axis(axes)
            new_nav_axes.append(longest)
            new_nav_shapes.append(
                longest.size
                if (ignore_axis is None or ignore_axis not in axes)
                else None
            )
        new_sig_axes = []
        new_sig_shapes = []
        for axes in zip_longest(
            *[s.axes_manager.signal_axes for s in args], fillvalue=None
        ):
            longest = _check_and_get_longest_axis(axes)
            new_sig_axes.append(longest)
            new_sig_shapes.append(
                longest.size
                if (ignore_axis is None or ignore_axis not in axes)
                else None
            )

        results = []
        new_axes = new_nav_axes[::-1] + new_sig_axes[::-1]
        new_data_shape = new_nav_shapes[::-1] + new_sig_shapes[::-1]
        for s in args:
            data = s._data_aligned_with_axes
            sam = s.axes_manager
            sdim_diff = len(new_sig_axes) - len(sam.signal_axes)
            while sdim_diff > 0:
                slices = (slice(None),) * sam.navigation_dimension
                slices += (None, Ellipsis)
                data = data[slices]
                sdim_diff -= 1
            thisshape = new_data_shape.copy()
            if ignore_axis is not None:
                _id = new_data_shape.index(None)
                newlen = data.shape[_id] if len(data.shape) > _id else 1
                thisshape[_id] = newlen
            thisshape = tuple(thisshape)
            if data.shape != thisshape:
                data = np.broadcast_to(data, thisshape)

            ns = s._deepcopy_with_new_data(data)
            ns.axes_manager._axes = [ax.copy() for ax in new_axes]
            ns.get_dimensions_from_data()
            results.append(ns.transpose(signal_axes=len(new_sig_axes)))
        return results


# functions for deciphering text

import re
import bisect
import collections
import copy

import numpy as np
import scipy

from wordbank import Token, Wordbank

vocab = None # modern dictionary word -> ID
inv_vocab = None # modern dictionary. ID -> word

def basic_substitute(ciphertext, wordbank):
    '''
    replace strings in ciphertext with known decryptions from wordbank

    ciphertext: string
    wordbank: instance of Wordbank
    '''

    token_list = tokenize_ciphertext(ciphertext)
    return ' '.join(token.plaintext for token in token_list)

# def apply_wordbank(token_list, wordbank):
#     for token in token_list:
#         if token in wordbank:
#             token.plaintext = wordbank[token]

# def load_vocab(filename):
#     '''
#     load a vocabulary.

#     input should be a sorted list
#     '''
#     _vocab = {} # word -> index
#     _inv_vocab = [] # index -> word
#     with open(filename) as f:
#         for idx, line in enumerate(f):
#             _vocab = {line: idx}
#             _inv_vocab.append(line)

#     _vocab['<unk>'] = len(_vocab)
#     _inv_vocab.append('<unk>')

#     return _vocab, _inv_vocab

# noisy channel model
def lm(token_list):
    '''
    read a context and return one-hot predictions over the output space

    output:
    '''

    # PLACEHOLDER UNTIL I GET A REAL LM
    # fake language model that simply uses the length. shorter words are better.

    score = 0

    # probability of each word. these are not history dependent
    likelihood = np.array([1 / len(word) for word in vocab.words])
    logprob = np.log(likelihood / likelihood.sum())

    # get score over history
    for token in token_list:
        score += logprob(vocab.lookup_word(token.plaintext))

    # add history score to score for each possible addition
    return score + logprob

def prior(token, wordbank):
    '''
    unigram log-probability of each token given its approximate location. returns an array of the same shape as lm
    model P(decoding|encoding) with a beta distribution (because it's normalized 0 to 1)
    '''

    left, right = wordbank.query_range(token)

    # find boundary words in modern vocab
    modern_left = bisect.bisect_left(modern_dict, left.plaintext)
    modern_right = bisect.bisect(modern_dict, right.plaintext)

    # set parameters of the beta distribution with mean where the word is relative to its anchor points
    mean = (int(token) - int(left)) / (int(right) - int(left))
    variance = 0.1
    n = mean * (1 - mean) / variance

    dist_params = {
        a: mean * n,
        b: (1 - mean) * n,
        loc: modern_left, # shift and scale to be about [0,1), and have variance 1
        scale: modern_right - modern_left,
    }
    if dist_params['a'] < 1 or dist_params['b'] < 1:
        raise ValueError('a, b must be > 1 to have a unimodal distribution')

    # probability of each word is the integral of the PDF from
    result = np.zeros(10000,)
    for word_idx in range(modern_left, modern_right):
        result[vocab[word]] = scipy.stats.beta.logcdf(word_idx + 1, **dist_params) - scipy.stats.beta.logcdf(word_idx, **dist_params)
    return result

Candidate = collections.namedtuple('Candidate', ['logprob', 'tokens'])

def decode(token_list, wordbank, beam_size=16):
    '''beam search decoding'''

    beam = [Candidate(logprob=0, tokens=[])]

    for token in token_list:
        new_candidates = []

        for beam_idx, candidate in enumerate(beam):
            if token.ciphertype is None or token.plaintext != '<unk>':
                # probably a hard-coded string. doesn't affect our model's probability
                candidate.tokens.append(token)
                new_candidates.append(candidate)

            elif token.ciphertype == 'dict':
                # prior is a beta distribution over the bounds established by the wordbank table

                lm_logprob = lm(candidate.tokens)
                prior_logprob = prior(wordbank, tokens)

                logprob = lm_logprob + prior_logprob

            elif token.ciphertype == 'table':
                # prior is a uniform distribution over all english words
                logprob = lm(candidate.tokens)

            # update beam
            if token.ciphertype == 'dict' or token.ciphertype == 'table':
                # get k best extensions to the current sequence and add them to the beam
                k = (beam_size - beam_idx)
                best_extensions = np.argpartition(logprob, len(vocab) - np.arange(k))[-k:] # np array of ints, shape: (k,)

                for ext_idx in best_extensions:
                    new_token = copy.copy(token)
                    new_token.plaintext = inv_vocab[ext_idx]
                    new_candidates.append(Candidate(logprob, candidate.tokens + [new_token]))

        beam = sorted(new_candidates)[-beam_size:] # todo: could save log(beam) time here by doing a partial sort only
        print(beam)

    print(beam)
    return beam[-1]

if __name__ == '__main__':
    wordbank = Wordbank()
    wordbank.load('wordbanks/wordbank.without.2880')

    # vocab, inv_vocab = load_vocab('dict.modern')
    ciphertext = '[556]^ 586[26]- Ferdinand 2 [678]^ 95 ? [1235]^ y ? [433]^ [79]^ [664]^'
    tokenized_ct = tokenize_ciphertext(ciphertext)

    # One hour baseline
    print(tokenized_ct)

    # Beam search experiments

    # print(decode(tokenize_ciphertext(ciphertext), wordbank))


from urllib.request import urlopen
import re

url = 'http://www.pythonchallenge.com/pc/def/ocr.html'
text = urlopen(url).read().decode()
symbols = re.findall("<!--(.*?)-->", text, re.DOTALL)[-1]
chars={}
for c in symbols:
    if c in chars:
        chars[c]+=1
    else:
        chars[c]=1
print(chars)


import numpy as np
from copy import deepcopy
from src.search.mcts.worker import MCTSWorker
from src.search.search import ParallelSearchAlgorithm
from torch.multiprocessing import Pipe

class MCTS(ParallelSearchAlgorithm):
    """ Monte Carlo Tree Search, with root parallelization."""

    def __init__(self, config):
        super().__init__(config)

        # Create workers
        pipes = [Pipe() for _ in range(self.num_workers)]
        self.channels = [p[0] for p in pipes]
        self.workers = [MCTSWorker(config, pipes[i][1]) for i in range(self.num_workers)]
        for w in self.workers:
            w.start()

    def search(self, state, iters=None):
        iters = iters if iters is not None else self.num_iters
        for c in self.channels:
            c.send({
                "command": "search",
                "state": deepcopy(state),
                "iters": iters,
            })
        msg = [c.recv() for c in self.channels]
        pi = np.mean([m["pi"] for m in msg], axis=0)
        qvals = np.mean([m["q"] for m in msg], axis=0)
        vals = np.mean([m["v"] for m in msg], axis=0)
        return {
            "pi": pi,
            "q": qvals,
            "v": vals,
        }


import os
import yaml
import numpy as np
import open3d as o3d
import random
from datasets.scannet200.owis_splits import KNOWN_CLASSES_LABELS,UNKNOWN_CLASSES_LABELS
from datasets.scannet200.scannet200_constants import CLASS_LABELS_200,VALID_CLASS_IDS_200
from copy import copy
# OUTPUT_2 = "output-scans/predictions-ours-2classes/"
label2id = {l:id for id,l in zip(VALID_CLASS_IDS_200, CLASS_LABELS_200)}
def generate_color_palette(x):
    palette = []

    while len(palette) < x:
        r, g, b = random.sample(range(256), 3)
        color = (r/255.0, g/255.0, b/255.0)
        if color[2] > 0:
            continue
        palette.append(color)

    return palette

def read_ply_file(file_path):
    # Read the ply file
    point_cloud = o3d.io.read_triangle_mesh(file_path)
    return point_cloud

model_name = "experiment_name"

for split in ['A','B','C']:
    for task in ['task1','task2','task3']:
        PATH = f"eval_output/{split}/{task}/instance_evaluation_{model_name}_0/decoder_-1/"
        label_database = "data/processed/scannet200/label_database.yaml"
        OUTPUT = f"output-scans/{model_name}/{split}/{task}/"
        OUTPUT_2 = f"output-scans/gt/{split}/{task}/"
        OUTPUT_3 = f"output-scans/input/{split}/{task}/"
        with open(label_database) as f:
            x = yaml.safe_load(f)

        # read label database
        cls_ids = x.keys()
        # from dict keys to list
        cls_ids = [int(i) for i in cls_ids]
        for file in os.listdir(PATH):
            pc = []
            pco = []
            SCENE = file.split(".")[0]
            # if SCENE in good_scenes:
            PATH_PLY = f"data/raw/scannet_test_segments/scans/{SCENE}/{SCENE}_vh_clean_2.labels.ply"
            PATH_PLY_1 = f"data/raw/scannet_test_segments/scans/{SCENE}/{SCENE}_vh_clean_2.ply"

            # read label database
            try:
                label = f"data/processed/scannet200/instance_gt/validation/{SCENE}.txt"
                # read label file
                with open(label) as f:
                    lines = f.readlines()
                unique_ids = np.unique(np.asarray(lines))
                unique_ids = [int(i[:-1]) for i in unique_ids]
                classes = [i//1000 for i in unique_ids if i !=0]
                # if task != 'task3':
                if UNKNOWN_CLASSES_LABELS[split][task] != None:
                    UNKNOWNS =  [label2id[i] for i in UNKNOWN_CLASSES_LABELS[split][task]]
                else:
                    UNKNOWNS = []

                instances = []
                if not os.path.exists(OUTPUT):
                    os.makedirs(OUTPUT)
                if not os.path.exists(OUTPUT_2):
                    os.makedirs(OUTPUT_2)
                if not os.path.exists(OUTPUT_3):
                    os.makedirs(OUTPUT_3)

                with open(f"{PATH}{SCENE}.txt") as f:
                    instances = f.readlines()
                # sort instances by probablity instances.split(" ")[2]
                instances = sorted(instances, key=lambda x: float(x.split(" ")[2]), reverse=False)

                pc = read_ply_file(PATH_PLY)
                pco = read_ply_file(PATH_PLY_1)
                points = np.asarray(pc.vertices)
                indices = []
                for instance in instances:
                    path = instance.split(" ")[0].strip()
                    cls = instance.split(" ")[1].strip()
                    # if(int(cls) not in cls_ids):
                    #     continue

                    with open(os.path.join(PATH, path), "r") as f:
                        mask = f.readlines()
                    for idx, point in enumerate(mask):
                        if(point == "1\n"):
                            indices.append(idx)
                            # change color of point based on its instance id
                            if(int(cls) == 0 or int(cls) == 1 or int(cls) == 3):
                                pc.vertex_colors[idx] = [0.5, 0.5, 0.5]
                            elif(int(cls) == 3000):
                                pc.vertex_colors[idx] = [31/255,119/255,180/255]
                            else:
                                pc.vertex_colors[idx] = [42/255,157/255,37/255]
                    indices = np.unique(np.asarray(indices))
                        # indices that are not in the list
                    left_indices = np.setdiff1d(np.arange(len(mask)), indices)
                    for index in left_indices:
                        pc.vertex_colors[index] = [0.5, 0.5, 0.5]
                    # save new point cloud
                    o3d.io.write_triangle_mesh(os.path.join(OUTPUT, SCENE + "-3d_owis.ply"), pc)

                    pc = read_ply_file(PATH_PLY)
                    points = np.asarray(pc.vertices)
                    for i in range(points.shape[0]):
                    # change color of point based on its instance id
                        if(int(lines[i][:-1])//1000 == 0 or int(lines[i][:-1])//1000 == 3 or int(lines[i][:-1])//1000 == 1): # make gray if 0 (nothing) or 3 (floor) or 1 (wall)
                            pc.vertex_colors[i] = [0.5, 0.5, 0.5]
                        elif(int(lines[i][:-1])//1000 in UNKNOWNS):
                            pc.vertex_colors[i] = [31/255,119/255,180/255]#[23/255,124/255,172/255]
                        else:
                            pc.vertex_colors[i] = [42/255,157/255,37/255] # green
                    o3d.io.write_triangle_mesh(os.path.join(OUTPUT_2, SCENE + "-gt.ply"), pc)
                    o3d.io.write_triangle_mesh(os.path.join(OUTPUT_3, SCENE + "-input.ply"), pco)
            except:
                pass


'''Some helper functions for PyTorch, including:
    - get_mean_and_std: calculate the mean and std value of dataset.
    - msr_init: net parameter initialization.
    - progress_bar: progress bar mimic xlua.progress.
'''
import os
import sys
import time
import math

import torch.nn as nn
import torch.nn.init as init

import numpy as np

def get_mean_and_std(dataset):
    '''Compute the mean and std value of dataset.'''
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=2)
    mean = torch.zeros(3)
    std = torch.zeros(3)
    print('==> Computing mean and std..')
    for inputs, targets in dataloader:
        for i in range(3):
            mean[i] += inputs[:,i,:,:].mean()
            std[i] += inputs[:,i,:,:].std()
    mean.div_(len(dataset))
    std.div_(len(dataset))
    return mean, std

def init_params(net):
    '''Init layer parameters.'''
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            init.kaiming_normal(m.weight, mode='fan_out')
            if m.bias:
                init.constant(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            init.constant(m.weight, 1)
            init.constant(m.bias, 0)
        elif isinstance(m, nn.Linear):
            init.normal(m.weight, std=1e-3)
            if m.bias:
                init.constant(m.bias, 0)

def computedistance2points(point1, point2):
    return math.sqrt((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)

def load_explore_traj(imgpath):
    if os.path.isdir(imgpath):
        imgs = []
        locations = []
        for name in os.listdir(imgpath):
            if name.endswith('.png'):
                imgs.append(name)
        imgs = sorted(imgs)
        for i in range(2,len(imgs),5):
            img  = imgs[i]
            img = img.split('.')[0]
            coord_x = float(img.split('_')[2])
            coord_y = float(img.split('_')[3])
            locations.append([coord_x, coord_y])
        return locations
    else:
        print("Not valid path")
        raise IOError

def load_navigation_traj_goal(path):
    position = []
    target = []
    with open(path,'r') as f:
        for cow in f.readlines():
            coord = list(map(float, cow.strip().split()))
            if len(target) == 0:
                target = [coord[2], coord[3]]
            position.append([coord[0], coord[1]])
    return position[:-1], target

_, term_width = os.popen('stty size', 'r').read().split()
term_width = int(term_width)

TOTAL_BAR_LENGTH = 65.
last_time = time.time()
begin_time = last_time
def progress_bar(current, total, msg=None):
    global last_time, begin_time
    if current == 0:
        begin_time = time.time()  # Reset for new bar.

    cur_len = int(TOTAL_BAR_LENGTH*current/total)
    rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1

    sys.stdout.write(' [')
    for i in range(cur_len):
        sys.stdout.write('=')
    sys.stdout.write('>')
    for i in range(rest_len):
        sys.stdout.write('.')
    sys.stdout.write(']')

    cur_time = time.time()
    step_time = cur_time - last_time
    last_time = cur_time
    tot_time = cur_time - begin_time

    L = []
    L.append('  Step: %s' % format_time(step_time))
    L.append(' | Tot: %s' % format_time(tot_time))
    if msg:
        L.append(' | ' + msg)

    msg = ''.join(L)
    sys.stdout.write(msg)
    for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3):
        sys.stdout.write(' ')

    # Go back to the center of the bar.
    for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2):
        sys.stdout.write('\b')
    sys.stdout.write(' %d/%d ' % (current+1, total))

    if current < total-1:
        sys.stdout.write('\r')
    else:
        sys.stdout.write('\n')
    sys.stdout.flush()

def format_time(seconds):
    days = int(seconds / 3600/24)
    seconds = seconds - days*3600*24
    hours = int(seconds / 3600)
    seconds = seconds - hours*3600
    minutes = int(seconds / 60)
    seconds = seconds - minutes*60
    secondsf = int(seconds)
    seconds = seconds - secondsf
    millis = int(seconds*1000)

    f = ''
    i = 1
    if days > 0:
        f += str(days) + 'D'
        i += 1
    if hours > 0 and i <= 2:
        f += str(hours) + 'h'
        i += 1
    if minutes > 0 and i <= 2:
        f += str(minutes) + 'm'
        i += 1
    if secondsf > 0 and i <= 2:
        f += str(secondsf) + 's'
        i += 1
    if millis > 0 and i <= 2:
        f += str(millis) + 'ms'
        i += 1
    if f == '':
        f = '0ms'
    return f

def adj2edges(adj):
    m = adj.shape[0]
    edges = []
    for i in range(m):
        for j in range(i):
            if adj[i, j] == 1:
                edges.append([i,j])
    return edges

def cos_sim(vector_a, vector_b):
    """
    Compute the cosine similarity
    :param vector_a: vector a
    :param vector_b: vector b
    :return: sim
    """
    vector_a = np.mat(vector_a)
    vector_b = np.mat(vector_b)
    num = float(vector_a * vector_b.T)
    denom = np.linalg.norm(vector_a) * np.linalg.norm(vector_b)
    cos = num / denom
    sim = 0.5 + 0.5 * cos
    return sim


from typing import Tuple, Union

class TrieNode:
    """
    A class to represent a node in a trie
    that contains only the pointers to the used characters
    (as opposed to the whole vocabulary array).

    ...

    Attributes
    ----------

    children : dict[str, TrieNode]
        pointers to the children nodes (for successive characters)
    """

    def __init__(self) -> None:
        self.children = {}

class Trie:

    def __init__(self, terminal: str = '$'):
        self.root = TrieNode()
        self.terminal = terminal

    def insert(self, word: str) -> None:
        """Inserts a word into the trie.

        Args:
            word (str): word to insert
        """
        if self.terminal in word:
            raise Exception(f'Explicit terminal ({self.terminal}) use is not allowed within words!')

        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()

            node = node.children[char]

        pointer_dict = node.children

        # proper terminating the word
        if self.terminal not in pointer_dict:
            pointer_dict[self.terminal] = None

    def search_prefix(self, prefix: str) -> Tuple[bool, Union[TrieNode, None]]:
        """Searches for given prefix in the trie

        Args:
            prefix (str): search term

        Returns:
            bool: True, if the prefix exists in trie
            TrieNode: - THE node that contains all suffixes of the prefix in trie,
                           if the prefix is contained
                      - None, if the prefix is not found
        """
        if self.terminal in prefix:
            raise Exception(f'Explicit terminal ({self.terminal}) use is not allowed within words!')

        node = self.root
        for char in prefix:
            if char in node.children:
                node = node.children[char]
            else:
                return False, None

        return True, node

    def search(self, word: str) -> bool:
        found, node = self.search_prefix(word)

        if not node:
            return False

        if self.terminal not in node.children:
            return False

        return True

test_trie1 = Trie()

vocabulary = ['cikla', 'cifar', 'aminokiselina', 'aminosulfat', 'amino']
[test_trie1.insert(word) for word in vocabulary]

search_terms = ['ami', 'amino', 'cikla', 'zirafa', 'cifara']
search_results = []
prefix_results = []
for term in search_terms:
    search_results.append(test_trie1.search(term))
    prefix_results.append(test_trie1.search_prefix(term)[0])

# test_trie1.insert('ami$') # causes exception
# test_trie1.search('ami$') # causes exception


#!/usr/bin/env python

from collections import defaultdict
from aoc import *
import re

bags_re = re.compile(r'(\d+) (.+?) bag')

def count_bags(contains: dict, have: str):
    cnt = 0
    for b, c in contains[have].items():
        cnt += c + c * count_bags(contains, b)
    return cnt

def solve(inp: str):
    contains = {}
    for line in inp.splitlines():
        outer, inner = line.split(" bags contain ", 2)
        contains[outer] = {bag: int(c) for c, bag in bags_re.findall(inner)}

    have = 'shiny gold'
    print(count_bags(contains, have))

# with open('test.txt', 'r') as f:
#     inp = f.read()
#     solve(inp)

with open('input.txt', 'r') as f:
    inp = f.read()
    solve(inp)


size = int(input())

matrix = []

for _ in range(size):
    matrix.append([int(x) for x in input().split()])

primary_diagonal = 0
secondary_diagonal = 0

for x in range(size):
    primary_diagonal += matrix[x][x]

for x in range(size):
    secondary_diagonal += matrix[x][size - x - 1]

print(abs(primary_diagonal - secondary_diagonal))


import sys, os
sys.path.append(os.getcwd())

from vocaloidDao import vocaloidVSQXResolver
import json

dataPath = os.path.join(os.getcwd(), 'VocaloidVSQXCollection')

with open(os.path.join(dataPath,"source.json"), 'r') as f:
    source = json.load(f)

fileList = [source["source-"+str(sourceIndex+1)]["file"] for sourceIndex in range(len(source))]

resolverList = [vocaloidVSQXResolver(os.path.join(dataPath, fileName)) for fileName in fileList]

import datetime
import multiprocessing as mp

def func(dataPath, resolver):
    return resolver.saveFormatedVocaloidData(resolver.getFormatedVocaloidDataInDict(), dataPath)

# # initialise the pool
num_cores = int(mp.cpu_count())
print("local computer has: " + str(num_cores) + " cores\n")
pool = mp.Pool(num_cores)

# # prepare processes not starting any process
processes = [pool.apply_async(func, args=(dataPath, resolver,)) for resolver in resolverList if not(
    os.path.exists(os.path.join(dataPath, resolver.resolverIdentifier[:-5]+"DataDict.csv")))]

start_t = datetime.datetime.now()
results = [p.get() for p in processes]
end_t = datetime.datetime.now()

elapsed_sec = (end_t - start_t).total_seconds()
print("Parallal computing takes " + "{:.2f}".format(elapsed_sec) + " seconds to finish.\n")


import os
import unittest
import sys

import pytest

from oarphpy_test import testutil
from oarphpy_test.testutil import skip_if_no_spark
from oarphpy_test.testutil import skip_if_tf_missing_or_broken

@skip_if_no_spark
def test_union_dfs():
  from oarphpy import spark as S
  with testutil.LocalSpark.sess() as spark:
    from pyspark.sql import Row
    df1 = spark.createDataFrame([Row(a=1, b=2.0)])
    df2 = spark.createDataFrame([Row(a=3, c='foo')])

    unioned = S.union_dfs(df1, df2)

    EXPECTED = """
         a    b     c
      0  1  2.0  None
      1  3  NaN   foo
    """
    actual = str(unioned.toPandas())
    def to_lines(str):
      lines = [line.strip() for line in str.split('\n')]
      return [l for l in lines if l]
    assert to_lines(EXPECTED) == to_lines(actual)

@skip_if_no_spark
def test_spark_selftest():
  assert testutil.LocalSpark.selftest(modname='oarphpy_test')

@skip_if_no_spark
def test_spark_has_lz4():
  import pyspark
  if pyspark.__version__.startswith('2.'):
    pytest.skip("Spark-2 environment doesn't have LZ4")

  TEST_TEMPDIR = testutil.test_tempdir('test_spark_has_lz4')

  from oarphpy import spark as S
  with testutil.LocalSpark.sess() as spark:
    from pyspark.sql import Row
    df = spark.createDataFrame([Row(a=1, b=2.0)] * 1000)
    df.write.parquet(path=os.path.join(TEST_TEMPDIR, 'out'), compression='lz4')

    dfr = spark.read.parquet(os.path.join(TEST_TEMPDIR, 'out'))
    assert dfr.count() == 1000

@skip_if_no_spark
def test_spark_has_zstd():
  import pyspark
  if pyspark.__version__.startswith('2.'):
    pytest.skip("Spark-2 environment doesn't have ZSTD")

  TEST_TEMPDIR = testutil.test_tempdir('test_spark_has_zstd')

  from oarphpy import spark as S
  with testutil.LocalSpark.sess() as spark:
    from pyspark.sql import Row
    df = spark.createDataFrame([Row(a=1, b=2.0)] * 1000)
    df.write.parquet(path=os.path.join(TEST_TEMPDIR, 'out'), compression='zstd')

    dfr = spark.read.parquet(os.path.join(TEST_TEMPDIR, 'out'))
    assert dfr.count() == 1000

@skip_if_no_spark
def test_spark_with_custom_library_script():
  # We must run this test in a subprocess in order for it to have an isolated
  # Spark session
  SCRIPT_PATH = testutil.get_fixture_path('test_spark_with_custom_library.py')

  from oarphpy import util
  out = util.run_cmd(sys.executable + ' ' + SCRIPT_PATH, collect=True)
  assert b'test_spark_with_custom_library success!' in out

@skip_if_no_spark
def test_spark_script_in_standalone_lib():
  # We must run this test in a subprocess in order for it to have an isolated
  # Spark session
  SCRIPT_PATH = testutil.get_fixture_path(
    'test_spark_script_in_standalone_lib.py')

  from oarphpy import util
  out = util.run_cmd(sys.executable + ' ' + SCRIPT_PATH, collect=True)
  assert b'test_spark_script_in_standalone_lib success!' in out

@skip_if_no_spark
def test_spark_with_custom_library_in_notebook():
  pytest.importorskip("jupyter")

  import re
  from oarphpy import util

  NB_PATH = testutil.get_fixture_path(
    'test_spark_ships_custom_library_in_notebook.ipynb')

  # We need to fork off nbconvert to run the test
  TEST_CMD = """
    cd /tmp &&
    jupyter-nbconvert \
        --ExecutePreprocessor.timeout=3600 \
        --to notebook --execute --output /tmp/out \
      {notebook_path}
  """.format(notebook_path=NB_PATH)
  out = util.run_cmd(TEST_CMD, collect=True)
  assert re.search('Writing .* bytes to /tmp/out.ipynb', out.decode())

@skip_if_no_spark
def test_cluster_get_info_smoke():
  with testutil.LocalSpark.sess() as spark:
    from oarphpy import spark as S
    S.cluster_get_info(spark)

@skip_if_no_spark
def test_pi():
  with testutil.LocalSpark.sess() as spark:
    from oarphpy import spark as S
    S.test_pi(spark)

@skip_if_no_spark
@skip_if_tf_missing_or_broken
def test_spark_tensorflow():
  from oarphpy import spark as S
  with testutil.LocalSpark.sess() as spark:
    S.test_tensorflow(spark)

@skip_if_no_spark
def test_spark_cpu_count():
  from oarphpy import spark as S
  import multiprocessing
  with testutil.LocalSpark.sess() as spark:
    assert multiprocessing.cpu_count() == S.cluster_cpu_count(spark)

@skip_if_no_spark
def test_spark_counters():
  from oarphpy import spark as S
  with testutil.LocalSpark.sess() as spark:
    counters = S.CounterCollection(spark)

    def run_counter(x, counters):

      # counters['iadd'] += 1
      counters['itally'] = 1
      counters.tally('my_tally', 2)
      counters.kv_tally('my_hist', key=str(x), value=1)

      return x

    task_rdd = spark.sparkContext.parallelize(range(10))

    with counters.log_progress():
      f = lambda x: run_counter(x, counters)
      final = task_rdd.map(f).sum()
      assert final == sum(range(10))

    assert counters['itally'] == 10
    assert "itally: 10" in str(counters)
    assert counters['my_tally'] == 20
    assert "my_tally: 20" in str(counters)
    assert \
      counters.get_kv_tally('my_hist') == dict((str(k), 1) for k in range(10))
    assert "my_hist: {" in str(counters)

@skip_if_no_spark
class TestArchiveRDD(unittest.TestCase):

  def _check_rdd(self, fw_rdd, fixture_strs):
    names = fw_rdd.map(lambda fw: fw.name).collect()
    expected_names = [s.decode('utf-8') for s in fixture_strs]
    assert sorted(names) == sorted(expected_names)

    values = fw_rdd.map(lambda fw: fw.data).collect()
    assert sorted(values) == sorted(fixture_strs)

  def test_archive_rdd_zip(self):
    TEST_TEMPDIR = testutil.test_tempdir('test_archive_rdd_zip')
    fixture_path = os.path.join(TEST_TEMPDIR, 'test.zip')

    # Create the fixture:
    # test.zip
    #  |- foo: "foo"
    #  |- bar: "bar"
    # ... an archive with a few files, where each file contains just a string
    #   that matches the name of the file in the archive
    ss = [b'foo', b'bar', b'baz']

    import zipfile
    with zipfile.ZipFile(fixture_path, mode='w') as z:
      for s in ss:
        z.writestr(s.decode('utf-8'), s)

    # Test Reading!
    with testutil.LocalSpark.sess() as spark:
      from oarphpy import spark as S
      fw_rdd = S.archive_rdd(spark, fixture_path)
      self._check_rdd(fw_rdd, ss)

  def test_archive_rdd_tar(self):
    TEST_TEMPDIR = testutil.test_tempdir('test_archive_rdd_tar')
    fixture_path = os.path.join(TEST_TEMPDIR, 'test.tar')

    # Create the fixture:
    # test.tar
    #  |- foo: "foo"
    #  |- bar: "bar"
    # ... an archive with a few files, where each file contains just a string
    #   that matches the name of the file in the archive
    ss = [b'foo', b'bar', b'baz']

    import tarfile
    with tarfile.open(fixture_path, mode='w') as t:
      for s in ss:
        from io import BytesIO
        buf = BytesIO()
        buf.write(s)
        buf.seek(0)

        buf.seek(0, os.SEEK_END)
        buf_len = buf.tell()
        buf.seek(0)

        info = tarfile.TarInfo(name=s.decode('utf-8'))
        info.size = buf_len

        t.addfile(tarinfo=info, fileobj=buf)

    # Test Reading!
    with testutil.LocalSpark.sess() as spark:
      from oarphpy import spark as S
      fw_rdd = S.archive_rdd(spark, fixture_path)
      self._check_rdd(fw_rdd, ss)

def _check_sample_in_expectation(df, n_per_category, expected, n_trials=10):
  import numpy as np
  import pandas as pd

  from oarphpy import spark as S

  df = df.cache()

  def _get_category_to_count(df):
    from collections import defaultdict
    rows = df.collect()

    category_to_count = defaultdict(int)
    for row in rows:
      category_to_count[row.val] += 1
    return category_to_count

  # NB: even with a fixed seed below, this test's behavior is dependent
  # on the number of system CPUs since Spark draws random numbers
  # concurrently.  To simulate a low-cpu system, use a Spark master local[1].
  # We do more trials to accomodate low-cpu systems (e.g. CircleCI).
  rows = [
    _get_category_to_count(
      S.get_balanced_sample(
        df, 'val',
        n_per_category=n_per_category,
        seed=100*s))
    for s in range(n_trials)
  ]
  pdf = pd.DataFrame(rows)
  pdf = pdf.fillna(0)

  ks = sorted(expected.keys())
  mu = pdf.mean()
  actual_arr = np.array([mu[k] for k in ks])
  expected_arr = np.array([expected[k] for k in ks])

  import numpy.testing as npt
  npt.assert_allclose(actual_arr, expected_arr, rtol=0.3)
    # NB: We can only test to about 30% accuracy with this few samples
    # If this fails, `try print(pdf)`

def _generate_rows():
  from pyspark.sql import Row

  VAL_TO_COUNT = {
    'a': 10,
    'b': 100,
    'c': 1000,
  }
  rows = []
  for val, count in VAL_TO_COUNT.items():
    for _ in range(count):
      i = len(rows)
      rows.append(Row(id=i, val=val))
  return rows

@skip_if_no_spark
def test_get_balanced_sample_simple():
  rows = _generate_rows()
  with testutil.LocalSpark.sess() as spark:
    df = spark.createDataFrame(rows)
    _check_sample_in_expectation(
      df, n_per_category=1, expected={'a': 1, 'b': 1, 'c': 1}, n_trials=30)

@skip_if_no_spark
def test_get_balanced_sample_sample_more_than_min():
  rows = _generate_rows()
  with testutil.LocalSpark.sess() as spark:
    df = spark.createDataFrame(rows)
    _check_sample_in_expectation(
      df, n_per_category=20, expected={'a': 10, 'b': 10, 'c': 10})

@skip_if_no_spark
def test_get_balanced_sample_sample_unrestricted():
  rows = _generate_rows()
  with testutil.LocalSpark.sess() as spark:
    df = spark.createDataFrame(rows)
    _check_sample_in_expectation(
      df, n_per_category=None, expected={'a': 10, 'b': 10, 'c': 10})

@skip_if_no_spark
@skip_if_tf_missing_or_broken
def test_spark_df_to_tf_dataset():
  import tensorflow as tf
  tf.compat.v1.disable_v2_behavior()

  from oarphpy.spark import spark_df_to_tf_dataset
  with testutil.LocalSpark.sess() as spark:

    import numpy as np
    import tensorflow as tf
    from pyspark.sql import Row

    def tf_dataset_to_list(ds):
      from oarphpy import util
      with util.tf_data_session(ds) as (sess, iter_dataset):
        return list(iter_dataset())

    df = spark.createDataFrame([
      Row(id='r1', x=1, y=[3., 4., 5.]),
      Row(id='r2', x=2, y=[6.]),
      Row(id='r3', x=3, y=[7., 8., 9.]),
    ])

    # Test empty
    ds = spark_df_to_tf_dataset(
            df.filter('x == False'), # Empty!,
            'id',
            spark_row_to_tf_element=lambda r: ('test',),
            tf_element_types=(tf.string,))
    assert tf_dataset_to_list(ds) == []

    # Test simple
    ds = spark_df_to_tf_dataset(
            df,
            'id',
            spark_row_to_tf_element=lambda r: (r.x,),
            tf_element_types=(tf.int64,))
    assert sorted(tf_dataset_to_list(ds)) == [(1,), (2,), (3,)]

    # Test Complex
    ds = spark_df_to_tf_dataset(
            df,
            'id',
            spark_row_to_tf_element=lambda r: (r.x, r.id, r.y),
            tf_element_types=(tf.int64, tf.string, tf.float64))
    expected = [
      (1, b'r1', np.array([3., 4., 5.])),
      (2, b'r2', np.array([6.])),
      (3, b'r3', np.array([7., 8., 9.])),
    ]
    items = list(zip(sorted(tf_dataset_to_list(ds)), sorted(expected)))
    for actual, exp in items:
      print('actual', actual,'exp', exp)
      assert len(actual) == len(exp)
      for i in range(len(actual)):
        np.testing.assert_array_equal(actual[i], exp[i])

################################################################################
### Test CloudpickeledCallable

def expensive_func():
  return "expensive_func"

def ExpensiveCallable(object):
  def __call__(self, x, y=5):
    return "ExpensiveCallable(%s, y=%s)" % (x, y)

@skip_if_no_spark
class TestCloudpickeledCallable(unittest.TestCase):

  def test_yay(self):
    from oarphpy.spark import CloudpickeledCallable
    f = CloudpickeledCallable(expensive_func)
    assert f() == expensive_func()
    assert str(f) == 'CloudpickeledCallable(_func_pyclass=oarphpy_test.test_spark.expensive_func)'

    cl = CloudpickeledCallable(lambda x: x * x)
    assert cl(2) == 4
    assert str(cl) == 'CloudpickeledCallable(_func_pyclass=oarphpy_test.test_spark.<lambda>)'

    import pickle
    b = pickle.dumps(f)

    decoded = pickle.loads(b)
    assert decoded() == expensive_func()
    assert str(decoded) == 'CloudpickeledCallable(_func_pyclass=oarphpy_test.test_spark.expensive_func)'
    assert decoded == f

################################################################################
### Test RowAdapter

## NB: these classes must be declared package-level (rather than test-scoped)
## for cloudpickle / Spark / RowAdapter to discover them properly

class Slotted(object):
  __slots__ = ('foo', 'bar', '_not_hidden')

  def __init__(self, **kwargs):
    # NB: ctor for convenience; SparkTypeAdapter does not require it
    for k in self.__slots__:
      setattr(self, k, kwargs.get(k))

  def __repr__(self):
    # For unit tests below
    return "Slotted(%s)" % ([(k, getattr(self, k)) for k in self.__slots__],)

class Unslotted(object):
  def __init__(self, **kwargs):
    for k, v in kwargs.items():
      setattr(self, k, v)

  def __repr__(self):
    # For unit tests below
    # Hide any 'hidden' attributes since `RowAdapter` will elide them
    attrs = dict(
              (k, v) for k, v in self.__dict__.items()
              if not k.startswith('__')
    )
    return "Unslotted(%s)" % str(sorted(attrs.items()))

  def __eq__(self, other):
    # Just for testing
    return repr(self) == repr(other)

# Tests for attrs-based classes are optional
try:
  import attr
  import numpy as np

  # NB: Need these classes defined package-level

  @attr.s(eq=True)
  class AttrsUnslotted(object):
    foo = attr.ib(default="moof")
    bar = attr.ib(default=5)

  @attr.s(slots=True, eq=True)
  class AttrsSlotted(object):
    foo = attr.ib(default="moof")
    bar = attr.ib(default=5)

except ImportError:
  pass

# Tests for dataclasses are optional
if sys.version_info[0] >= 3:
  from oarphpy_test.dataclass_obj import *

# We use Row a lot in the tests below; tests are skipped w/out Spark
try:
  from pyspark.sql import Row
except Exception:
  pass

def _select_distinct(df, col):
  return list(set(r[0] for r in df.select(col).collect()))

@skip_if_no_spark
class TestRowAdapter(unittest.TestCase):

  ## Tests

  def test_nonadapted_input(self):
    from oarphpy.spark import RowAdapter

    # RowAdapter leaves bare data input unchanged
    BARE_VALUES = True, 1, 1.0, "moof", bytes(b"moof")
    for datum in BARE_VALUES:
      assert RowAdapter.to_row(datum) == datum
      assert RowAdapter.from_row(datum) == datum

  def test_pesky_numpy(self):
    import numpy as np
    from oarphpy.spark import RowAdapter

    # RowAdapter translates pesky numpy-boxed numbers ...
    assert RowAdapter.to_row(Row(x=np.float32(1.))) == Row(x=1.)

    # ... but only one way! In practice, just don't save boxed numbers in rows.
    assert RowAdapter.from_row(Row(x=np.float32(1.))) == Row(x=np.float32(1.))

  def test_python_basic_values(self):

    ## Columns with basic python types get translated as-is; FMI see Spary's
    ## core type mappings: https://github.com/apache/spark/blob/master/python/pyspark/sql/types.py#L875
    row_expected_schema = [
      (Row(x=True),                 [('x', 'boolean')]),
      (Row(x=1),                    [('x', 'bigint')]),
      (Row(x=1.0),                  [('x', 'double')]),
      (Row(x="moof"),               [('x', 'string')]),
      (Row(x=bytearray(b"moof")),   [('x', 'binary')]),
      (Row(x=None),                 [('x', 'void')]),
    ]

    if self._is_spark_2x():
      row_expected_schema[-1] = (
        (Row(x=None),                 [('x', 'null')]))

    self._check_raw_adaption([(r, r) for r, s in row_expected_schema])

    for row, expected_schema in row_expected_schema:
      schema = self._check_schema([row], expected_schema)
      VOID_TYPE = 'void' if not self._is_spark_2x() else 'null'
      if all(t != VOID_TYPE for colname, t in expected_schema):
        # In most cases, data can be written as expected
        self._check_serialization([row])
      else:
        # NB: None / null / void can't be written
        with pytest.raises(Exception) as excinfo:
          self._check_serialization([row], schema=schema)
        assert ((
          "Parquet data source does not support %s data type" % VOID_TYPE)
          in str(excinfo.value))

  def test_python_basic_containers_adaption(self):

    ## RowAdapters translates basic python containers recursively
    rows = [
      Row(x=[]),
      Row(x=tuple()),
      Row(x={}),

      Row(x=[1, 2, 3]),
      Row(x=(True, False)),
      Row(x={'a': 1, 'b': 2}),

      Row(x=[{'a': 1}, {'c': 2}]),

      Row(x=[Row(a={'b': Row(c=None)})]),
        # NB: as we'll see below, RowAdapter translates
        # `dict` as "map" and `Row` as "struct"
    ]
    self._check_raw_adaption([(r, r) for r in rows])

  def test_python_basic_empty_containers(self):

    # NOTE!!  Empty containers can't have their types auto-deduced!!
    rows = [
      Row(id=0, l=[]),
      Row(id=1, l=[]),
    ]
    with pytest.raises(ValueError) as excinfo:
      self._check_serialization(rows)
    assert "Some of types cannot be determined" in str(excinfo.value)

    # RowAdapter can deduce the type, but `array<null>` can't be serialized
    # (same with map<null, null>)
    schema = self._check_schema(
                      rows,
                      [('id', 'bigint'),
                      ('l', 'array<void>')
                        if not self._is_spark_2x() else ('l', 'array<null>')
                      ])
    with pytest.raises(Exception) as excinfo:
      self._check_serialization(rows, schema=schema)
    if self._is_spark_2x():
      assert ("Parquet data source does not support array<null> data type"
        in str(excinfo.value))
    else:
      assert ("Parquet data source does not support array<void> data type"
        in str(excinfo.value))

    # WORKAROUND: If you compute a schema based upon a prototype row, then you
    # can use that schema to write empty containers
    PROTOTYPE_ROW = Row(id=0, l=[0])
    schema = self._check_schema(
                      [PROTOTYPE_ROW],
                      [('id', 'bigint'), ('l', 'array<bigint>')])
    self._check_serialization(rows, schema=schema)

  def test_python_basic_nonempty_containers(self):

    # Data with non-empty containers enjoys auto schema deduction from Spark
    rows = [
      Row(id=0, l=[1],    d={'a': 'foo'}),
      Row(id=1, l=[2, 3], d={'b': 'bar'}),
      Row(id=2, l=[],     d={}),    # Some rows with empty containers are OK
      Row(id=3, l=None,   d=None)   # Null data is OK too
    ]
    self._check_serialization(rows)
    self._check_schema(
            rows,
            [('id', 'bigint'),
             ('l', 'array<bigint>'),
             ('d', 'map<string,string>')
            ])

  def test_basic_structs(self):

    # Spark treats `Row` like a 'struct' rather than a map; here's how you can
    # use raw `Row` instances:
    rows = [
      Row(id=0, x=Row(shape='square', area=1)),
      Row(id=1, x=Row(shape='circle', area=2)),
    ]
    self._check_serialization(rows)
    self._check_schema(
            rows,
            [('id', 'bigint'),
             ('x',  'struct<shape:string,area:bigint>'),
            ])

  def test_built_in_unslotted(self):

    # Spark actually has built-in support for adapting plain Python objects,
    # but ...
    rows = [
      Row(id=0, x=Unslotted(v=4)),
      Row(id=1, x=Unslotted(v=5)),
    ]
    df = self._check_serialization(rows, do_adaption=False)

    # ... Spark deserializes 'objects' as generic `Row` instances.
    df_rows = df.collect()
    assert sorted(df_rows) == sorted([
      Row(id=0, v=Row(v=4)), Row(id=1, x=Row(v=5))
    ])

  def test_built_in_dataclasses(self):
    pytest.importorskip('dataclasses')

    # Spark also has built-in support for Python dataclasses, but ...
    rows = [
      Row(id=0, x=DataclassObj(x='foo', y=5.)),
      Row(id=1, x=DataclassObj(x='bar', y=6.)),
    ]
    df = self._check_serialization(rows, do_adaption=False)

    # ... Spark deserializes 'dataclasses' as generic `Row` instances.
    df_rows = df.collect()
    assert sorted(df_rows) == sorted([
      Row(id=0, v=Row(x='foo', y=5.)), Row(id=1, x=Row(x='bar', y=6.))
    ])

  def test_built_in_attrs(self):
    pytest.importorskip('attr')

    # Spark also has built-in support for attrs-based objects, but ...
    rows = [
      Row(id=0, x=AttrsUnslotted()),
      Row(id=1, x=AttrsUnslotted(foo="foom")),
    ]
    df = self._check_serialization(rows, do_adaption=False)

    # ... Spark deserializes 'attrs' objects as generic `Row` instances.
    df_rows = df.collect()
    assert sorted(df_rows) == sorted([
      Row(id=0, x=Row(bar=5, foo='moof')), Row(id=1, x=Row(bar=5, foo='foom'))
    ])

  def test_built_in_numpy(self):
    import numpy as np

    # Spark can NOT handle complex objects like numpy arrays
    rows = [
      Row(id=0, x=np.array([1, 2, 3])),
      Row(id=1, x=np.array([4, 5, 6])),
    ]
    with pytest.raises(TypeError) as excinfo:
      self._check_serialization(rows, do_adaption=False)

    if self._is_spark_2x():
      assert "not supported type: <class 'numpy.ndarray'>" in str(excinfo.value)
    else:
      # pyspark>=3.3.1
      assert "Unable to infer the type of the field x" in str(excinfo.value)

  def test_built_in_slotted(self):

    # Spark does NOT have built-in support for *slotted* objects
    rows = [
      Row(id=0, x=Slotted(foo=5, bar="abc", _not_hidden=1)),
      Row(id=1, x=Slotted(foo=7, bar="cba", _not_hidden=3)),
    ]
    with pytest.raises(TypeError) as excinfo:
      self._check_serialization(rows, do_adaption=False)

    if self._is_spark_2x():
      assert \
        ("not supported type: <class 'oarphpy_test.test_spark.Slotted'>"
        in str(excinfo.value))
    else:
      # pyspark>=3.3.1
      assert "Unable to infer the type of the field x" in str(excinfo.value)

  def test_built_in_contained_udt(self):
    from pyspark.ml.linalg import DenseVector
    from pyspark.sql.types import StructType

    # As of Spark 3.x, Spark DOES support lists, maps, and structs of UDTs ...
    assert hasattr(DenseVector, '__UDT__')
    rows = [
      Row(id=0, x=[DenseVector([-1.0, -1.0])]),
      Row(id=0, x=[DenseVector([-2.0, -2.0])]),
    ]
    df = self._check_serialization(rows, do_adaption=False)

    # ... and will reconstruct python objects when querying data ...
    assert df.collect() == rows

    # But there are several drawbacks:
    #  1) You must manually write the schema for your UDT using Spark's type
    #       objects:
    assert type(DenseVector.__UDT__.sqlType()) == StructType
    #  2) For ndarrays, Spark ML only supports double-valued structures and
    #        offers no packed encoding for large arrays; compare with
    #        RowAdapter's Tensor below and see also:
    #          https://github.com/apache/spark/blob/6c805470a7e8d1f44747dc64c2e49ebd302f9ba4/python/pyspark/ml/linalg/__init__.py#L144
    values_field = [
      f for f in DenseVector.__UDT__.sqlType().fields if f.name == 'values'
    ][0]

    if self._is_spark_2x():
      assert str(values_field.dataType) == 'ArrayType(DoubleType,false)'
    else:
      assert str(values_field.dataType) == 'ArrayType(DoubleType(), False)'

  def test_rowadapter_unslotted(self):

    # RowAdapter will deserialize and re-create Python objects, using the
    # `Unslotted` class defined at runtime.
    rows = [
      Row(id=0, x=Unslotted(v=4)),
      Row(id=1, x=Unslotted(v=5)),
    ]
    df = self._check_serialization(rows)

    # NB: The check above also checks that `Unslotted` instances are created
    # and checks equality via __eq__ or pprint.pformat().  But just to make
    # things visible in this test:
    from oarphpy.spark import RowAdapter
    decoded = sorted(RowAdapter.from_row(r) for r in df.collect())
    assert [r.x for r in decoded] == [Unslotted(v=4), Unslotted(v=5)]

    # RowAdapter records the (full) class name in the table as the
    # `__pyclass__` attribute of each value of the `x` column.
    assert _select_distinct(df, 'x.__pyclass__') == [
      'oarphpy_test.test_spark.Unslotted']

    # RowAdapter encodes objects as structs (even though in Python objects are
    # very dict-like).
    self._check_schema(
            rows,
            [('id', 'bigint'),
             ('x',  'struct<__pyclass__:string,v:bigint>'),
            ])

  def test_rowadapter_dataclasses(self):
    pytest.importorskip('dataclasses')

    rows = [
      Row(id=0, x=DataclassObj(x='foo', y=5.)),
      Row(id=1, x=DataclassObj(x='bar', y=6.)),
    ]
    df = self._check_serialization(rows)

    assert _select_distinct(df, 'x.__pyclass__') == [
      'oarphpy_test.dataclass_obj.DataclassObj']

    self._check_schema(
            rows,
            [('id', 'bigint'),
             ('x',  'struct<__pyclass__:string,x:string,y:double>'),
            ])

  def test_rowadapter_attrs(self):
    pytest.importorskip('attr')

    rows = [
      Row(id=0, x=AttrsUnslotted()),
      Row(id=1, x=AttrsUnslotted(foo="foom")),
    ]
    df = self._check_serialization(rows)
    assert _select_distinct(df, 'x.__pyclass__') == [
      'oarphpy_test.test_spark.AttrsUnslotted']
    self._check_schema(
            rows,
            [('id', 'bigint'),
             ('x',  'struct<__pyclass__:string,foo:string,bar:bigint>'),
            ])

    rows = [
      Row(id=0, x=AttrsSlotted()),
      Row(id=1, x=AttrsSlotted(foo="foom")),
    ]
    df = self._check_serialization(rows)
    assert _select_distinct(df, 'x.__pyclass__') == [
      'oarphpy_test.test_spark.AttrsSlotted']

    self._check_schema(
            rows,
            [('id', 'bigint'),
             ('x',  'struct<__pyclass__:string,foo:string,bar:bigint>'),
            ])

  def test_rowadapter_numpy_unpacked(self):
    import numpy as np

    # RowAdapter translates Numpy arrays to a oarphy Tensor object that affords
    # SQL-based inspection for small arrays (and uses a more efficient row- or
    # column-major packed encoding for large arrays; see next test)
    rows = [
      Row(id=0, x=np.array([1, 2, 3])),
      Row(id=1, x=np.array([4, 5, 6])),
    ]
    df = self._check_serialization(rows)
    assert _select_distinct(df, 'x.__pyclass__') == [
              'oarphpy.spark.Tensor']
    EXPECTED = """
        id                                                     x
    0   0  (oarphpy.spark.Tensor, [3], int64, C, [1, 2, 3], [])
    1   1  (oarphpy.spark.Tensor, [3], int64, C, [4, 5, 6], [])
    """
    self._pandas_compare_str(df.orderBy('id').toPandas(), EXPECTED)

    # _check_serialization() verifies that the data gets decoded as numpy
    # arrays, but just to make things visible in this test:
    from oarphpy.spark import RowAdapter
    decoded = [RowAdapter.from_row(r) for r in df.orderBy('id').collect()]
    np.testing.assert_equal(decoded[0].x, np.array([1, 2, 3]))
    np.testing.assert_equal(decoded[1].x, np.array([4, 5, 6]))

  def test_rowadapter_numpy_packed(self):
    import numpy as np
    from oarphpy.spark import TENSOR_AUTO_PACK_MIN_KBYTES

    N = int(TENSOR_AUTO_PACK_MIN_KBYTES * (2**10) / np.dtype(int).itemsize) + 1
    expect_packed = np.reshape(np.array(range(2 * N)), (2, N))

    rows = [
      Row(id=0, x=expect_packed),
      Row(id=1, x=(expect_packed + 1)),
    ]
    df = self._check_serialization(rows)
    assert _select_distinct(df, 'x.__pyclass__') == [
              'oarphpy.spark.Tensor']

    # Verify that we actually have a column of packed values
    bin_data = df.select('*').first().x.values_packed
    assert len(bin_data) == expect_packed.size * expect_packed.dtype.itemsize
      # For ints, usually 8 bytes per int * 2 * N

    # _check_serialization() verifies that the data gets decoded as numpy
    # arrays, but just to make things visible in this test:
    from oarphpy.spark import RowAdapter
    decoded = [RowAdapter.from_row(r) for r in df.orderBy('id').collect()]
    np.testing.assert_equal(decoded[0].x, expect_packed)
    np.testing.assert_equal(decoded[1].x, expect_packed + 1)

  def test_rowadapter_cloudpickled_callable(self):
    from oarphpy.spark import CloudpickeledCallable

    def moof():
      return 'moof'
    cc_moof = CloudpickeledCallable(moof)
    assert cc_moof() == moof()

    cc_empty = CloudpickeledCallable()
    assert cc_empty == CloudpickeledCallable.empty()
    with pytest.raises(Exception):
      cc_empty()

    rows = [
      Row(id=0, f=cc_empty),
      Row(id=1, f=cc_moof),
    ]
    df = self._check_serialization(rows)
    assert _select_distinct(df, 'f.__pyclass__') == [
              'oarphpy.spark.CloudpickeledCallableData']

    # Ensuring invoking the functions still works
    from oarphpy.spark import RowAdapter
    decoded = [RowAdapter.from_row(r) for r in df.orderBy('id').collect()]
    with pytest.raises(Exception):
      decoded[0].f()
    assert decoded[1].f() == moof()

  def test_rowadapter_slotted(self):

    # RowAdapter will treat slotted objects like normal (unslotted) objects.
    # The main difference is that if the slotted type changes and _loses_
    # slots, RowAdapter won't try to setattr() on missing slots.
    rows = [
      Row(id=0, x=Slotted(foo=5, bar="abc", _not_hidden=1)),
      Row(id=1, x=Slotted(foo=7, bar="cba", _not_hidden=3)),
    ]
    df = self._check_serialization(rows)
    assert _select_distinct(df, 'x.__pyclass__') == [
      'oarphpy_test.test_spark.Slotted']

    # RowAdapter encodes objects as structs (even though in Python objects are
    # very dict-like).
    self._check_schema(
            rows,
            [('id', 'bigint'),
             ('x',  'struct<__pyclass__:string,foo:bigint,bar:string,_not_hidden:bigint>'),
            ])

  def test_rowadapter_contained_objects(self):

      # RowAdapter will process containers of translatable objects recursively
      s1 = Slotted(foo=5, bar="abc", _not_hidden=1)
      s2 = Slotted(foo=7, bar="cba", _not_hidden=3)
      rows = [
        Row(id=0, x={'a': [{'b': [s1]}]}),
        Row(id=1, x={'a': [{'b': [s2]}]}),
      ]
      df = self._check_serialization(rows)

      # The resulting type is very complicated ...
      self._check_schema(
              rows,
              [('id', 'bigint'),
              ('x',  'map<string,array<map<string,array<struct<__pyclass__:string,foo:bigint,bar:string,_not_hidden:bigint>>>>>'),
              ])

      # ... but not hard to query!
      assert 0 ==  df.select('*').where("x.a[0].b[0].foo = 5").first().id
      assert 1 ==  df.select('*').where("x.a[0].b[0].bar = 'cba'").first().id

  def test_rowadapter_complex(self):
    from oarphpy.spark import RowAdapter

    # A large-ish example that covers the above cases in aggregate
    rows = [
      Row(
        id=1,
        np_number=np.float32(1.),
        a=np.array([1]),
        b={
          'foo': np.array( [ [1] ], dtype=np.uint8)
        },
        c=[
          np.array([[[1.]], [[2.]], [[3.]]])
        ],
        d=Slotted(foo=5, bar="abc", _not_hidden=1),
        e=[Slotted(foo=6, bar="def", _not_hidden=1)],
        f=Unslotted(meow=4, _not_hidden=1, __hidden=2),
        g=Unslotted(), # Intentionally empty; adapter should set nothing
        h=Row(i=1, j=2),
      ),

      # Include a mostly empty row below to exercise Spark type validation.
      # Spark will ensure the row below and row above have the same schema;
      # note that `None` (or 'null') is only allowed for Struct / Row types.
      Row(
        id=2,
        np_number=np.float32(2.),
        a=np.array([]),
        b={},
        c=[],
        d=None,
        e=[],
        f=None,
        g=None,
        h=Row(i=3, j=3),
      ),
    ]

    df = self._check_serialization(rows)
    EXPECTED_ALL = """
                                                                                0                                                1
      id                                                                            1                                                2
      np_number                                                                   1.0                                              2.0
      a                                (oarphpy.spark.Tensor, [1], int64, C, [1], [])  (oarphpy.spark.Tensor, [0], float64, C, [], [])
      b              {'foo': ('oarphpy.spark.Tensor', [1, 1], 'uint8', 'C', [1], [])}                                               {}
      c          [(oarphpy.spark.Tensor, [3, 1, 1], float64, C, [1.0, 2.0, 3.0], [])]                                               []
      d                                  (oarphpy_test.test_spark.Slotted, 5, abc, 1)                                             None
      e                                [(oarphpy_test.test_spark.Slotted, 6, def, 1)]                                               []
      f                                     (oarphpy_test.test_spark.Unslotted, 4, 1)                                             None
      g                                          (oarphpy_test.test_spark.Unslotted,)                                             None
      h                                                                        (1, 2)                                           (3, 3)
      """

    # DEPRECATED: pyspark 2.x is deprecated
    import pyspark
    if pyspark.__version__.startswith('2.'):
      EXPECTED_ALL = """
                                                                                0                                                1
        id                                                                            1                                                2
        np_number                                                                     1                                                2
        a                                (oarphpy.spark.Tensor, [1], int64, C, [1], [])  (oarphpy.spark.Tensor, [0], float64, C, [], [])
        b              {'foo': ('oarphpy.spark.Tensor', [1, 1], 'uint8', 'C', [1], [])}                                               {}
        c          [(oarphpy.spark.Tensor, [3, 1, 1], float64, C, [1.0, 2.0, 3.0], [])]                                               []
        d                                  (oarphpy_test.test_spark.Slotted, 5, abc, 1)                                             None
        e                                [(oarphpy_test.test_spark.Slotted, 6, def, 1)]                                               []
        f                                     (oarphpy_test.test_spark.Unslotted, 4, 1)                                             None
        g                                          (oarphpy_test.test_spark.Unslotted,)                                             None
        h                                                                        (1, 2)                                           (3, 3)
        """

    self._pandas_compare_str(df.orderBy('id').toPandas().T, EXPECTED_ALL)

    # Test Schema Deduction
    mostly_empty = Row(
      id=2,
      np_number=None,
      a=None,
      b={},
      c=[],
      d=None,
      e=[],
      f=None,
      g=None,
      h=None,
    )
    mostly_empty_adapted = RowAdapter.to_row(mostly_empty)

    # Spark can't deduce schema from the empty-ish row ...
    with pytest.raises(ValueError) as excinfo:
      self._check_serialization([mostly_empty_adapted], do_adaption=False)
    assert "Some of types cannot be determined" in str(excinfo.value)

    # ... but this works if we tell it the schema!
    schema = RowAdapter.to_schema(rows[0])
    self._check_serialization([mostly_empty_adapted], schema=schema)

    # Let's check that RowAdapter schema deduction works as expected
    EXPECTED_SCHEMA = [
      ('id', 'bigint'),
      ('np_number', 'double'),
      ('a', 'struct<__pyclass__:string,shape:array<bigint>,dtype:string,order:string,values:array<bigint>,values_packed:binary>'),
      ('b', 'map<string,struct<__pyclass__:string,shape:array<bigint>,dtype:string,order:string,values:array<bigint>,values_packed:binary>>'),
      ('c', 'array<struct<__pyclass__:string,shape:array<bigint>,dtype:string,order:string,values:array<double>,values_packed:binary>>'),
      ('d', 'struct<__pyclass__:string,foo:bigint,bar:string,_not_hidden:bigint>'),
      ('e', 'array<struct<__pyclass__:string,foo:bigint,bar:string,_not_hidden:bigint>>'),
      ('f', 'struct<__pyclass__:string,meow:bigint,_not_hidden:bigint>'),
      ('g', 'struct<__pyclass__:string>'),
      ('h', 'struct<i:bigint,j:bigint>'),
    ]
    self._check_schema(rows, EXPECTED_SCHEMA)

  ## Test Support

  def _is_spark_2x(self):
    if not hasattr(self, '_is_spark_2x_cache'):
      import pyspark
      self._is_spark_2x_cache = pyspark.__version__.startswith('2.')
    return self._is_spark_2x_cache

  def _check_raw_adaption(self, raw_expected):
    from oarphpy.spark import RowAdapter

    for raw_data, expected_row in raw_expected:
      actual_row = RowAdapter.to_row(raw_data)
      assert actual_row == expected_row

      actual_data = RowAdapter.from_row(expected_row)
      assert actual_data == raw_data

  def _check_schema(self, rows, expected_schema):
    from oarphpy.spark import RowAdapter
    schema = RowAdapter.to_schema(rows[0])
    adapted_rows = [RowAdapter.to_row(r) for r in rows]
    with testutil.LocalSpark.sess() as spark:
      df = spark.createDataFrame(
        adapted_rows, schema=schema, verifySchema=False)
        # verifySchema is expensive and improperly errors on mostly empty rows

      if self._is_spark_2x():
        # Spark 2.x returns schema values in a different order, so we do a more
        # flexible test
        def tokenize(s):
          import re
          return sorted(re.split('[<>,]+', s))
        actual = dict((col, tokenize(s)) for col, s in df.dtypes)
        expected = dict((col, tokenize(s)) for col, s in expected_schema)
        assert actual == expected
      else:
        # Tests are written for Spark 3.x
        assert df.dtypes == expected_schema

    return schema

  def _check_serialization(self, rows, schema=None, do_adaption=True):
    import inspect
    from oarphpy import util
    from oarphpy.spark import RowAdapter

    test_name = inspect.stack()[1][3]

    TEST_TEMPDIR = testutil.test_tempdir('TestRowAdapter.' + test_name)

    if do_adaption:
      adapted_rows = [RowAdapter.to_row(r) for r in rows]
    else:
      adapted_rows = rows

    with testutil.LocalSpark.sess() as spark:
      if schema:
        df = spark.createDataFrame(
          adapted_rows, schema=schema, verifySchema=False)
            # verifySchema is expensive and improperly errors on mostly
            # empty rows
      else:
        df = spark.createDataFrame(adapted_rows)
          # Automatically samples rows to get schema
      outpath = os.path.join(TEST_TEMPDIR, 'rowdata_%s' % test_name)
      df.write.parquet(outpath)

      df2 = spark.read.parquet(outpath)
      decoded_wrapped_rows = df2.collect()

      if do_adaption:
        decoded_rows = [
          RowAdapter.from_row(row)
          for row in decoded_wrapped_rows
        ]
        # We can't do assert sorted(rows) == sorted(decoded_rows)
        # because numpy syntatic sugar breaks __eq__, so use pprint,
        # which is safe for our tests
        import pprint
        def sorted_row_str(rowz):
          if self._is_spark_2x():
            # Spark 2.x has non-stable sorting semantics for Row
            if len(rowz) > 1:
              rowz = sorted(rowz, key=lambda r: r.id)
            return pprint.pformat(rowz)
          else:
            return pprint.pformat(sorted(rowz))
        assert sorted_row_str(rows) == sorted_row_str(decoded_rows)

      return df

  def _pandas_compare_str(self, pdf, expected):
    def cleaned(s):
      lines = [l.strip() for l in s.split('\n')]
      lines = [l for l in lines if l]
      if self._is_spark_2x():
        # Spark 2.x returns table rows and array values in different orders,
        # so we do a more flexible test
        def tokenize(s):
          import re
          return sorted(re.split('[\(\),]+', s))
        lines = sorted(' '.join(tokenize(l)) for l in lines)
      return '\n'.join(lines)
    actual = pdf.to_string()
    assert cleaned(actual) == cleaned(expected), \
      "\n\nExpected %s \n\nActual %s\n\n" % (expected, actual)


"""
cbook即为cookbook，是一些小工具组成的库

"""
from matplotlib.cbook import CallbackRegistry

callbacks = CallbackRegistry()
sum = lambda x, y: print(f'{x}+{y}={x + y}')
mul = lambda x, y: print(f"{x} * {y}={x * y}")
id_sum = callbacks.connect("sum", sum)
id_mul = callbacks.connect("mul", mul)
callbacks.process('sum', 3, 4)
callbacks.process("mul", 5, 6)
callbacks.disconnect(id_sum)
callbacks.process("sum", 7, 8)


"""Common methods we reuse."""
import collections
import os
import shutil
import stat
import tempfile as tmp
from contextlib import contextmanager
from subprocess import check_output

def sort_tag(tag):
    """Sort the list by tag."""
    list_length = len(tag)

    # Length 8 is list -l, 7 is df, 5 is list
    if list_length == 9:
        _tag = tag[4]
    elif list_length == 7:
        _tag = tag[6]
    elif list_length == 6:
        _tag = tag[3]
    else:
        _tag = tag[1]

    _sort = _tag.rsplit('_', 1)

    # We want to sort tags that have been created with count > 1. But not
    # foo_bar
    if len(_sort) > 1 and _sort[1].isdigit():
        return _sort[0], int(_sort[1])
    else:
        return _tag, 0

def sort_release(releases, split=False):
    """
    Sort the list by RELEASE, if split is true it's expecting full
    datasets.
    """
    r_dict = {}
    release_list = []

    if split:
        for rel in releases:
            rel, r_type = rel.properties["mountpoint"].value.rsplit("/")[
                -1].split("-")

            if len(rel) > 2:
                rel = float(rel)

            r_dict[rel] = r_type
    else:
        for release in releases:
            try:
                release, r_type = release.split("-")

                if len(release) > 2:
                    release = float(release)

                r_dict[release] = r_type
            except ValueError:
                pass

    ordered_r_dict = collections.OrderedDict(sorted(r_dict.items()))
    index = 0

    for r, t in ordered_r_dict.items():
        if split:
            release_list.insert(index, ["{}-{}".format(r, t)])
            index += 1
        else:
            release_list.insert(index, "{}-{}".format(r, t))
            index += 1

    return release_list

# Cyrille Pontvieux on StackOverflow
def copytree(src, dst, symlinks=False, ignore=None):
    """Copies a tree and overwrites."""
    if not os.path.exists(dst):
        os.makedirs(dst)
        shutil.copystat(src, dst)
    lst = os.listdir(src)
    if ignore:
        excl = ignore(src, lst)
        lst = [x for x in lst if x not in excl]
    for item in lst:
        s = os.path.join(src, item)
        d = os.path.join(dst, item)
        if symlinks and os.path.islink(s):
            if os.path.lexists(d):
                os.remove(d)
            os.symlink(os.readlink(s), d)
            try:
                st = os.lstat(s)
                mode = stat.S_IMODE(st.st_mode)
                os.lchmod(d, mode)
            except:
                pass  # lchmod not available
        elif os.path.isdir(s):
            copytree(s, d, symlinks, ignore)
        else:
            shutil.copy2(s, d)

# http://stackoverflow.com/questions/2333872/atomic-writing-to-file-with-python
@contextmanager
def tempfile(suffix='', dir=None):
    """
    Context for temporary file.

    Will find a free temporary filename upon entering
    and will try to delete the file on leaving, even in case of an exception.

    Parameters
    ----------
    suffix : string
        optional file suffix
    dir : string
        optional directory to save temporary file in
    """

    tf = tmp.NamedTemporaryFile(delete=False, suffix=suffix, dir=dir)
    tf.file.close()
    try:
        yield tf.name
    finally:
        try:
            os.remove(tf.name)
        except OSError as e:
            if e.errno == 2:
                pass
            else:
                raise

@contextmanager
def open_atomic(filepath, *args, **kwargs):
    """
    Open temporary file object that atomically moves to destination upon
    exiting.

    Allows reading and writing to and from the same filename.

    The file will not be moved to destination in case of an exception.

    Parameters
    ----------
    filepath : string
        the file path to be opened
    fsync : bool
        whether to force write the file to disk
    *args : mixed
        Any valid arguments for :code:`open`
    **kwargs : mixed
        Any valid keyword arguments for :code:`open`
    """
    fsync = kwargs.get('fsync', False)

    with tempfile(dir=os.path.dirname(os.path.abspath(filepath))) as tmppath:
        with open(tmppath, *args, **kwargs) as file:
            try:
                yield file
            finally:
                if fsync:
                    file.flush()
                    os.fsync(file.fileno())
        os.rename(tmppath, filepath)
        os.chmod(filepath, 0o644)

def indent_lines(message):
    """This will indent all lines except the first by 7 spaces. """
    indented = message.replace("\n", "\n       ").rstrip()

    return indented

def get_nested_key(_dict, keys=[]):
    """Gets a nested key from a dictionary."""
    key = keys.pop(0)

    if len(keys) == 0:
        return _dict[key]

    return get_nested_key(_dict[key], keys)

def checkoutput(*args, **kwargs):
    """Just a wrapper to return utf-8 from check_output"""
    out = check_output(*args, **kwargs).decode("utf-8")

    return out


def factorial(valor):
    bandera = 1
    primera = 1
    segunda = 2
    while bandera < valor:
        primera = primera * segunda
        segunda = segunda + 1
        bandera = bandera + 1
    return(primera)

def swap(lista, indiceUno, indiceDos):
    auxi = lista[indiceUno]
    lista[indiceUno] = lista[indiceDos]
    lista[indiceDos] = auxi


class Game:
    def __init__(self, rows=14, cols=18, mines=40):
        self.rows = rows
        self.cols = cols
        self.mines = mines
        self.board = [[Cell() for _ in range(self.cols)]
                      for _ in range(self.rows)]
        self.set_mine()
        self.set_position()

    def set_mine(self):
        x = 0
        while x <= self.mines:
            for i in range(self.mines//10):
                for j in range(self.mines//4):
                    self.board[i][j].mine = True
                    self.board[i][j].neighbors = 1
                    x += 1

    def set_position(self):
        for i in range(self.rows):
            for j in range(self.cols):
                self.board[i][j].row = i
                self.board[i][j].col = j

    # all(i.neighbors == sum(b[x][y].mine for x, y in [(i.row+x, i.col+y) for x, y in [(x, y) for x in range(-1, 2) for y in range(-1, 2) if x or y]] if 0 <= x < r and 0 <= y < c) for i in lst)

class Cell:

    def __init__(self):
        self.open = False
        self.mine = False
        self.neighbors = 0
        self.flag = False
        self.row = 0
        self.col = 0

game = Game()
# game.rows ➞ 14
# game.cols ➞ 18
# game.mines ➞ 40

# g1 = (Game(), 14, 18, 40)
# g2 = (Game(20, 20, 20), 20, 20, 20)
# g3 = (Game(20, 20, 300), 20, 20, 300)
# g4 = (Game(25, 25, 200), 25, 25, 200)
# print('---------------------------------')

# for g, r, c, m in [g1, g2, g3, g4]:
#     b = g.board
#     lst = [cell for row in b for cell in row]

#     print('Tests for Game({}, {}, {})'.format(r, c, m))
#     Test.assert_equals(g.rows, r, 'Attribute rows')
#     Test.assert_equals(g.cols, c, 'Attribute cols')
#     Test.assert_equals(g.mines, m, 'Attribute mines')
#     Test.assert_equals(len(b), r, 'Board height')
#     Test.assert_equals(len(b[0]), c, 'Board width')
#     Test.assert_equals(sum(type(i) is Cell for i in lst),
#                        r*c, 'Number of Cells')
#     Test.assert_equals(sum(i.mine for i in lst), m, 'Number of mines')
#     Test.assert_equals(all((b[i][j].row, b[i][j].col) == (
#         i, j) for i in range(r) for j in range(c)), True, 'Positions match')
#     Test.assert_equals(all(not i.flag or i.open for i in lst),
#                        True, 'Attributes flag and open')
#     Test.assert_equals(all(i.neighbors == sum(b[x][y].mine for x, y in [(i.row+x, i.col+y) for x, y in [(
#         x, y) for x in range(-1, 2) for y in range(-1, 2) if x or y]] if 0 <= x < r and 0 <= y < c) for i in lst), True, 'Neighbors')
#     print('---------------------------------')


import string
from enum import Enum

lower_alphabet_list = list(string.ascii_lowercase)
upper_alphabet_list = list(string.ascii_uppercase)
number_list = list(string.digits)

class EncryptionMode(Enum):
    ENC = 'ENCRYPT'
    DEC = 'DECRYPT'

def vigenere_crypt_key_range(key: str) -> int:
    key_ord = ord(key)

    if key in lower_alphabet_list:
        key_ord = key_ord - ord('a')
    elif key in upper_alphabet_list:
        key_ord = key_ord - ord('A')
    elif key in number_list:
        key_ord = key_ord - ord('0')

    return key_ord

def vigenere_crypt_pass(ch:str, key:int) -> str:
    target_list = None
    ch_ord = ord(ch)

    pos = 0
    pos_range = 26

    if ch in lower_alphabet_list:
        pos = ch_ord - ord('a')
        target_list = lower_alphabet_list
    elif ch in upper_alphabet_list:
        pos = ch_ord - ord('A')
        target_list = upper_alphabet_list
    elif ch in number_list:
        pos = ch_ord - ord('0')
        target_list = number_list
        pos_range = 10

    return target_list[(pos + key) % pos_range]

def vigenere_encrypt_decrypt(text: str, key: str, mode: EncryptionMode) -> str:
    """
    비제네르 암호를 이용하여 암호화 혹은 복호화를 수행하는 암호 알고리즘
    :param text: 암호화할 문자열
    :param key: 암호화에 사용할 key의 배열
    :param mode: 암호화할 지 복호화할 지 구분하기 위한 값
    :return: 비제네르 암호를 이용한 암호문 혹은 복호화된 문자열
    """

    key_len     = len(key)
    enc_mode    = 1 if mode == EncryptionMode.ENC else -1
    result_list = []

    for i in range(0, len(text)):
        mod_key = vigenere_crypt_key_range(key[i % key_len]) * enc_mode
        result_list.append(vigenere_crypt_pass(text[i], mod_key))

    return str(''.join(result_list))


import socket

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
server_address1 = ('localhost', 1234)
message = 'trying to connect'

try:
	sent = sock.sendto(message.encode(), server_address1)
	port, server = sock.recvfrom(4096)
	s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	server_address = ('localhost', int(port.decode()))
	alist = ' '.join([input() for x in range(5)])
	sent = s.sendto(alist.encode(), server_address)
	data, server = s.recvfrom(4096)
	print(data.decode())

finally:
	sock.close()