"""
Yatao (An) Bian <yatao.bian@gmail.com>
yataobian.com
May 13, 2019.
"""
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os, sys
sys.path.append('./' )
import datetime
import pickle
from collections import defaultdict
from absl import flags
from absl import app

from functions import utils
from itertools import product, combinations
from functions.exp_specs import  (solver_names, amazon_categories, amazon_ns,
plot_opt, data_name, all_folds, result_path_elbo, result_path_paelbo)

flags.DEFINE_boolean('debug', True, 'Whether it is in debug mode.')
flags.DEFINE_integer('problem_id', 1, 'Options: 1: ELBO, 2: PA-ELBO.')
flags.DEFINE_string('mode', 'run', 'Options: run: run algorithms; stats: get experimental statistics.')
FLAGS = flags.FLAGS
FLAGS(sys.argv)

class exp_result():
    '''dummy class to act as struct'''
    pass

class exp_config():
    '''dummy class to act as struct'''
    pass

np.random.seed(99)
nm_latent_dims = 3 # number of latent dimensionalities in FLID
Ds = [2, 3, 10]

# Set solvers
if 1 == FLAGS.problem_id:  # elbo
    # solver 4 always needed
    # 2 shrunken FW does not plot trajectory
    one_epoch_solvers = [0, 3, 11]
    multi_epoch_solvers = [4,5,6,9,10,12]
    solver_list = one_epoch_solvers + multi_epoch_solvers

    # path to store generated figures
    data_fig_path = result_path_elbo
    extra_epoch = 5   #real number of epochs is 1 + extra_epoch

    if FLAGS.debug:
        folds = [3]
        folds_iter = [folds]
        category_id = 10
        category_iter = [category_id]

    # run all categories and folds
    else:
        folds_iter = list(combinations (all_folds, 1))
        category_iter =np.arange( len ( amazon_ns))

    """Control whether to plot the true log-partition or not. The true ELBO is from
    Tschiatschek, S., Djolonga, J., and Krause, A.
    Learning probabilistic submodular diversity models via noise contrastive estimation.
    In Proc. International Conference on Artificial Intelligence and Statistics (AISTATS), 2016.
    """
    MIT_TRUE_EVIDENCE = 1

elif 2 == FLAGS.problem_id:  # pa-elbo
    one_epoch_solvers = [0, 3, 11]
    multi_epoch_solvers = [4,5,6,9,10,12]
    solver_list = one_epoch_solvers + multi_epoch_solvers

    data_fig_path = result_path_paelbo

    extra_epoch = 9  # real # epoch is 1 + extra_epoch

    if FLAGS.debug:
        folds = (1,2)
        folds_iter = [folds]
        category_id = 11
        category_iter = [category_id]
    else:
        folds_iter = list (combinations (all_folds, 2) )
        category_iter = np.arange( len ( amazon_ns))

    # there is no true log-partition in this setting.
    MIT_TRUE_EVIDENCE = 0

else:
    pass

nm_solver_names = len(solver_names)
print('solver_list: ', solver_list, 'nm_solver_names:', nm_solver_names)

def run_algorithms():
    for category_id in category_iter:
        print('processing category id: {} ...'.format(category_id))

        for folds in folds_iter:
            category = amazon_categories[category_id]
            n = amazon_ns[category_id]
            co_order = []  #  coordinate order
            for i in range(1 + extra_epoch):
                co_order.append(np.random.permutation(n))

            nm_exps = 1
            Dstepsize = 1

            # used for traj. plotting
            exp_id = 0; param_id = 0
            subfix = '-cate-' + category + '-n-'+ str(n) +\
                 '-dataid-' + str(FLAGS.problem_id) + '-folds-' + str(folds)

            results = defaultdict(exp_result)   # use a dict

            for idx_exp  in range(nm_exps): # repeated experiments
                for i in range(nm_latent_dims):  #i:  param idx
                    Dtmp = Ds[i]
                    # load data
                    [f, grad, gradi, param] = utils.load_flid_data(FLAGS.problem_id,
                                                                   Dtmp,
                                                                   n,
                                                                   data_fig_path,
                                                                   category,
                                                                   folds)

                    print('data: %s, category: %s,  n: %d, D: %d, folds: %s \n'\
                       % ( data_name[FLAGS.problem_id], category,  param.n, Dtmp, str(folds)))

                    for solver_id in solver_list:
                        [opt_x, opt_f, fs, margs, runtime] =\
                         utils.launch_solver(f, grad, gradi,
                            param, solver_id, extra_epoch,
                            co_order=co_order)

                        results[idx_exp, i, solver_id].opt_f = opt_f
                        results[idx_exp, i, solver_id].fs = fs
                        results[idx_exp, i, solver_id].opt_x = opt_x
                        results[idx_exp, i, solver_id].margs = margs
                        results[idx_exp, i, solver_id].runtime = runtime
                        results[idx_exp, i, solver_id].logZ = param.logZ

            # record experimental settings
            expconfig = exp_config()
            expconfig.n = n; expconfig.Ds=Ds; expconfig.nm_exps=nm_exps
            expconfig.Dstepsize=Dstepsize; expconfig.K=nm_latent_dims
            expconfig.solver_list=solver_list; expconfig.exp_id =exp_id
            expconfig.nm_names = nm_solver_names
            expconfig.data_id = FLAGS.problem_id
            expconfig.time = datetime.date.today()

            try:
                original_umask = os.umask(0)
                os.makedirs(data_fig_path, exist_ok=True, mode=0o777)
            finally:
                os.umask(original_umask)
            utils.supermakedirs(data_fig_path, mode=0o777)

            file_name = data_fig_path + 'mfi-' + data_name[FLAGS.problem_id] \
                + subfix

            pickle.dump([results, expconfig], open(file_name +'.pkl', 'wb'))

            # ------- do the plot ---------
            if FLAGS.problem_id == 1:
                ylabel = 'ELBO'
            elif FLAGS.problem_id == 2:
                ylabel = 'PA-ELBO'
            else:
                raise ValueError('Unknown problem id: {}!'.format(FLAGS.problem_id) )

            tight_height = 0.86
            sns.set( )
            sns.set_style('whitegrid')
            from matplotlib import rc
            SMALL_SIZE = 6
            MEDIUM_SIZE = 6
            BIGGER_SIZE = 6
            rc('legend', markerscale=1)
            rc('lines', linewidth=0.4)
            rc('font', size=SMALL_SIZE)
            rc('axes', titlesize=SMALL_SIZE)
            rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
            rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
            rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
            rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

            opt_fs = np.zeros( [nm_exps, nm_latent_dims, nm_solver_names] )
            for id in range(nm_exps):
                for i in range(nm_latent_dims):
                    for t in solver_list:
                        opt_fs[id, i, t] = results[id, i, t].opt_f

            opt_mean = np.mean(opt_fs, axis=0).squeeze()
            opt_std = np.squeeze( np.std(opt_fs, axis=0))

            opt_max = np.max(np.max(opt_mean[:, solver_list] ))
            opt_min = np.min(np.min(opt_mean[:, solver_list] ))

            # plot returned function value
            fig_name = data_name[FLAGS.problem_id] + '-' + subfix
            plt.close('all')
            fig_scale = 1
            fWidth = 4*fig_scale
            fHeight = 4*fig_scale

            # by default markersize is 12
            plt_lw = 1; markersize = 5

            def plot_elbo_values(solver_list, subfix='', ylabel='ELBO'):
              # elbo value with legend
              hFig = plt.figure(figsize=(fWidth, fHeight*tight_height ))

              hands = []
              for t in solver_list[1:]:
                  hi = plt.plot(np.arange(len(Ds)), opt_mean[:, t], plot_opt[t],\
                                lw=plt_lw, markersize=markersize)
                  hands.append(hi[0])
                  plt.errorbar(np.arange(len(Ds)), opt_mean[:, t], yerr=opt_std[:, t],\
                                fmt=plot_opt[t],lw=plt_lw, markersize=markersize)

              t = solver_list[0]
              hi = plt.plot(np.arange(len(Ds)), opt_mean[:, t], plot_opt[t], \
                            lw=plt_lw, markersize=markersize)
              hands.append(hi[0])
              plt.errorbar(np.arange(len(Ds)), opt_mean[:, t], yerr=opt_std[:, t], \
                           fmt=plot_opt[t],lw=plt_lw, markersize=markersize)

    #           # plot true logZ
    #           # using plot option 8 in exp_specs
              if MIT_TRUE_EVIDENCE and subfix=='multiepoch':
                  solver_id = solver_list[0]
                  plt.plot( np.arange(len(Ds)) , \
                            [results[exp_id, k, solver_id].logZ for k in range(nm_latent_dims) ], plot_opt[8],\
                                markersize=markersize, lw=plt_lw)
              figure_part = 0.74
              plt.tight_layout(rect= [0.08,0.04, 1, figure_part] )
              newlist = solver_list[1:] + solver_list[0:1]

              plt.legend(hands, [solver_names[i] for i in newlist], \
                  bbox_to_anchor=(-0.09, 1, 1.1, 0.35), loc='lower left', \
                     mode='expand', fontsize='large',\
                     ncol =2)

          #    plt.axis([Ds[0], Ds[-1],   opt_min, opt_max])
          #    plt.axis(  )
              plt.xticks(np.arange(len(Ds)), Ds)
              plt.xlabel('D')
              plt.ylabel(ylabel)
              sns.despine()
              plt.savefig(data_fig_path + fig_name+ subfix+'.pdf', format='pdf', \
                          transparent=False)

              # ------ plot ELBO value without legend
              hFig = plt.figure(figsize=(fWidth, fHeight*tight_height ))
              hands = []
              for t in solver_list[1:]:
                  hi = plt.plot(np.arange(len(Ds)), opt_mean[:, t], plot_opt[t],\
                                lw=plt_lw, markersize=markersize)
                  hands.append(hi[0])
                  plt.errorbar(np.arange(len(Ds)), opt_mean[:, t], yerr=opt_std[:, t],\
                                fmt=plot_opt[t],lw=plt_lw, markersize=markersize)

              t = solver_list[0]
              hi = plt.plot(np.arange(len(Ds)), opt_mean[:, t], plot_opt[t], \
                            lw=plt_lw, markersize=markersize)
              hands.append(hi[0])
              plt.errorbar(np.arange(len(Ds) ), opt_mean[:, t], yerr=opt_std[:, t], \
                           fmt=plot_opt[t],lw=plt_lw, markersize=markersize)

              # plot true logZ
              # plot option 8
              if MIT_TRUE_EVIDENCE and subfix=='multiepoch':
                  plt.plot(np.arange(len(Ds)) , \
                            [results[exp_id, k, solver_id].logZ for k in range(nm_latent_dims) ], plot_opt[8],\
                                markersize=markersize, lw=plt_lw)

              # plt.tight_layout(rect= [0.08,0.04, 1, 1])
              plt.xticks(np.arange(len(Ds)), Ds)
              plt.xlabel('D')
              plt.ylabel(ylabel)
              sns.despine()
              plt.savefig(data_fig_path + fig_name+ subfix+'nolegend.pdf', format='pdf', \
                          transparent=False)
              plt.draw()

            plot_elbo_values(one_epoch_solvers, subfix='1epoch', ylabel=ylabel)
            plot_elbo_values(multi_epoch_solvers, subfix='multiepoch', ylabel=ylabel)

            def plot_trajectories(solver_list, is_multi_epoch=True, ylabel='ELBO'):
              #    --------------- plot trajectories of the solvers-----
              solvers = solver_list.copy()
              if is_multi_epoch:
                subfix2 = '_multiepoch'
                # not plot traj. for solver 2, shrunken FW
                if 2 in solvers:
                  solvers.remove(2)
              else:
                subfix2 = '_1epoch'
              if is_multi_epoch:
                iters = np.arange(0, n*(extra_epoch+1) + 1)
              else:
                iters = np.arange(0, n+1)

              markersize = 1
              for param_id_ in  range(len(Ds)):
                  fig_name = data_name[FLAGS.problem_id] + subfix + \
                      '-traj-expid-' + str(exp_id) + '-D-' + str(Ds[param_id_] )

                  hFig = plt.figure(figsize=(fWidth, fHeight*tight_height ))
                  hands = []
                  for t in solvers[1:]:
                      tr = results[exp_id, param_id_, t].fs
                      if (isinstance(tr, list) and 1==len(tr) ):
                          tr = np.ones( [ len(iters), 1] ) *tr
                      elif not isinstance(tr, list) and  1==tr.size:
                          tr = np.ones( [len(iters), 1] ) *tr

                      hi = plt.plot(iters, tr, plot_opt[t],\
                          linewidth=plt_lw, markersize= markersize)
                      hands.append(hi[0])

                  t = solvers[0]
                  tr = results[exp_id, param_id_, t].fs
                  if 1 ==len(tr):
                      tr = np.ones( [1, len(iters)] ) *tr

                  hi =plt.plot(iters, tr, plot_opt[t],\
                      lw=plt_lw, markersize= markersize)
                  hands.append(hi[0])

                  if is_multi_epoch:
                    # plot  1-epoch vertical line
                    plt.axvline(x=n, color='c', linestyle='--', lw = plt_lw*0.5)

                  # plot true logZ
                  # plot style 8
                  if MIT_TRUE_EVIDENCE and is_multi_epoch:
                      plt.axhline(results[exp_id, param_id_, 4].logZ, color='y', linestyle='-', marker='d',\
                                    markersize=markersize, lw=plt_lw)

                  newlist = solvers[1:] + solvers[0:1]
                  plt.legend(hands, [solver_names[i] for i in newlist], \
                       loc='best', \
                         mode=None, fontsize='x-large',\
                         ncol =1)

                  plt.xlabel(r'Iterations. $D={}$'.format(Ds[param_id_]));
                  plt.ylabel(ylabel)
                  if is_multi_epoch:
                      if MIT_TRUE_EVIDENCE:
                          maxf = results[exp_id, param_id_, 4].logZ
                          plt.ylim(ymin=-2, ymax = 0.4 + maxf)
                      else:
                          plt.ylim(ymin=-2, ymax=0.4+opt_max)

                  sns.despine()
                  plt.savefig( data_fig_path+fig_name+subfix2+'.pdf', fmt='pdf', \
                               transparent=False)

            plot_trajectories(multi_epoch_solvers, is_multi_epoch=True, ylabel=ylabel)
            plot_trajectories(one_epoch_solvers, is_multi_epoch=False, ylabel=ylabel)

"""
Compute experimental statistics for one-epoch solvers.
The statistics is used to generate Table 1 in the paper.
"""
def compute_exp_stats():

    if FLAGS.problem_id == 1:
        # cate x D x solver x fold
        optf_1epoch = np.zeros([len(amazon_categories), len(Ds), nm_solver_names, len(all_folds)])
    else:  # PA-ELBO
       # cate x D x solver x fold_pair
        optf_1epoch = np.zeros([len(amazon_categories), len(Ds), nm_solver_names, len(folds_iter)])

    for category_id in category_iter:
        for folds_id, folds in  enumerate(folds_iter):
            print(folds[0])
            category = amazon_categories[category_id]
            n = amazon_ns[category_id]

            nm_solver = len(solver_list)
            nm_exps = 1
            Dstepsize = 1

            exp_id=0; param_id = 0
            subfix = '-cate-' + category + '-n-'+ str(n) +\
                  '-dataid-' + str(FLAGS.problem_id) + '-folds-'+str(folds)

            file_name = data_fig_path + 'mfi-' + data_name[FLAGS.problem_id] + subfix + '.pkl'
            if os.path.getsize(file_name) > 0:
                results, expconfig = pickle.load(open( file_name, 'rb' ))

            for d in range(nm_latent_dims):
                for s in one_epoch_solvers:
                    optf_1epoch[category_id, d, s, folds_id] = results[0, d, s].opt_f

    if FLAGS.problem_id == 1:
        file_name = os.path.join(data_fig_path, 'optf_1epoch')
    else:
        file_name = os.path.join(data_fig_path, 'optf_1epoch_pa')

    pickle.dump(optf_1epoch, open(file_name +'.pkl', 'wb'))

def main(_):
    if FLAGS.mode == 'run':
        run_algorithms()
    elif FLAGS.mode == 'stats':
        compute_exp_stats()
    else:
        raise ValueError('Unknown mode: {}'.format(FLAGS.mode))

if __name__ == '__main__':
    app.run(main)


#!/usr/bin/env python3

import sys
import time

import collections
from enum import Enum

from abc import ABC, abstractmethod

from PyQt5.QtCore import Qt

__device_comm_write = sys.__stdout__
__device_comm_read = sys.__stdin__

class Device(ABC):

    _DEVICE_TYPE_MAP = {}

    def __init__(self, device_path='', parent=None):

        if isinstance(device_path, Device):
            parent = device_path.parent
            device_path = device_path.device_path

        self.__device_path = device_path
        self.__parent = parent

        if self.sendMessage('get-info is-device-group') == 'yes':

            children_count = int(self.sendMessage('device-count'))

            if device_path == '':
                dev_path_prefix = ':'
            else:
                dev_path_prefix = device_path

            children = []
            for i in range(children_count):
                child_type = self.sendMessage(f'{i}: device-type')
                child_class = self._DEVICE_TYPE_MAP.get(child_type, Device)
                children.append(child_class(
                    device_path=f'{device_path}{i}:', parent=self))

            self.__children = tuple(children)
        else:
            self.__children = None

        self.__device_type = self.sendMessage('device-type')
        self.__device_desc = self.sendMessage('device-desc')
        self.__device_name = self.sendMessage(
            'get-info device-name-in-group')
        if self.__device_name == '<<null>>':
            self.__device_name = self.__device_type

    @property
    def device_path(self):
        return self.__device_path

    @property
    def parent(self):
        return self.__parent

    @property
    def children(self):
        return self.__children

    @property
    def type_(self):
        return self.__device_type

    @property
    def name(self):
        return self.__device_name

    @property
    def description(self):
        return self.__device_desc

    def __get_repr(self, depth=0):

        spaces = '  '*depth
        prefix = spaces + self.__device_name

        if self.__children is None:
            return prefix

        children_repr = (child.__get_repr(depth+1) for child in self.__children)
        return prefix + ': {\n' + ',\n'.join(children_repr) + f'\n{spaces}}}'

    def __repr__(self):
        return self.__get_repr()

    def sendMessage(self, message):
        return send(self.__device_path + message)

class Engine(Device):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.__intensity = None

    @property
    def intensity(self):
        if self.__intensity is None:
            self.__intensity = self.sendMessage('get-property intensity')

        return self.__intensity

    @intensity.setter
    def intensity(self, value):
        self.__intensity = None
        self.sendMessage(f'set-property intensity {value}')

Device._DEVICE_TYPE_MAP['engine'] = Engine
Device._DEVICE_TYPE_MAP['linear-engine'] = Engine

class Sensor(Device):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        try:
            self.__reading_time = float(self.sendMessage('reading-time'))
            self.__max_offset = float(self.sendMessage('max-offset'))
            self.__max_error = float(
                self.sendMessage('max-error')) - self.__max_offset
        except ValueError:
            self.__reading_time = self.__max_offset = self.__max_error = None

    @property
    def reading_time(self):
        return self.__reading_time

    @property
    def max_error(self):
        return self.__max_error

    @property
    def max_offset(self):
        return self.__max_offset

    def read(self, subdevice=None):
        if subdevice is None:
            msg = 'read'
        else:
            msg = f'{subdevice}:read'

        try:
            return float(self.sendMessage(msg))
        except ValueError:
            return float('NaN')

Device._DEVICE_TYPE_MAP['position-sensor'] = Sensor
Device._DEVICE_TYPE_MAP['angle-sensor'] = Sensor
Device._DEVICE_TYPE_MAP['speed-sensor'] = Sensor
Device._DEVICE_TYPE_MAP['ang-speed-sensor'] = Sensor
Device._DEVICE_TYPE_MAP['line-dist-sensor'] = Sensor

SensorInfo = collections.namedtuple('SensorInfo', (
    'reading_time', 'max_error', 'max_offset', 'estimated_offset'))

class TextInputDevice(Device):

    @abstractmethod
    def read(self, size=-1):
        pass

    @abstractmethod
    def readline(self):
        pass

    def flush(self):
        pass

class SimpleKeyboardInputDevice(TextInputDevice):

    def __init__(self, device_path='', parent=None):
        super().__init__(device_path=device_path, parent=parent)

        self.__buffer = None

    def read(self, size=-1):

        if self.__buffer is None:
            self.__buffer = self.__read()
        elif size < 0 or len(self.__buffer) < size:
            self.__buffer += self.__read()

        if size < 0 or len(self.__buffer) < size:
            read_content = self.__buffer
            self.__buffer = None
            return read_content

        read_content = self.__buffer[:size]
        self.__buffer = self.__buffer[size:]

        return read_content

    def readline(self):
        return self.read()

    def flush(self):
        pass

    def _returnToBuffer(self, text):
        if self.__buffer is None:
            self.__buffer = text
        else:
            self.__buffer = text + self.__buffer

    def __read(self):
        return self.sendMessage('get')

Device._DEVICE_TYPE_MAP['keyboard'] = SimpleKeyboardInputDevice

class TranslatedKeyboardInputDevice(SimpleKeyboardInputDevice):

    Key = collections.namedtuple('Key', (
        'char', 'code', 'modifiers', 'shift', 'control'))

    def read(self, size=-1, end_at=None):
        content = super().read(size)
        key_list = []

        for i in range(0, len(content) - 9, 10):

            key_content = int(content[i: i+10], 16)

            key_modifiers = key_content >> 8
            key_code = key_content & 0xffffffff

            shift_up = key_modifiers & Qt.ShiftModifier

            key_char = None
            if key_code >= Qt.Key_Space and key_code <= Qt.Key_AsciiTilde:
                ascii_code = key_code
                if not shift_up:
                    if key_code >= Qt.Key_A and key_code <= Qt.Key_Z:
                        ascii_code = key_code + ord('a') - ord('A')
                key_char = chr(ascii_code)

            elif key_code == Qt.Key_Return:
                key_char = '\n'

            key_list.append(self.Key(
                char=key_char,
                code=key_code,
                modifiers=key_modifiers,
                shift=shift_up,
                control=key_modifiers & Qt.ControlModifier
            ))

            if end_at is not None:
                if key_char == end_at or key_code == end_at:
                    self._returnToBuffer(content[i+10:])
                    return key_list

        return key_list

    def readline(self):
        return self.read(end_at='\n')

class TextOutputDevice(Device):

    @abstractmethod
    def write(self, text):
        pass

    def flush(self):
        pass

_SimpleConsoleOutputDevice_WriteMode = Enum(
    'WriteMode', ('Normal', 'Push', 'Static'))

class SimpleConsoleOutputDevice(TextOutputDevice):

    WriteMode = _SimpleConsoleOutputDevice_WriteMode

    def __init__(self, device_path='', parent=None):
        super().__init__(device_path=device_path, parent=parent)

        self.__message_buffer = ''
        self.__current_message_mode = \
            SimpleConsoleOutputDevice.WriteMode.Normal;

    def write(self, text, mode=_SimpleConsoleOutputDevice_WriteMode.Normal):
        if mode is not self.__current_message_mode:
            if self.__message_buffer:
                self.flush()
            self.__current_message_mode = mode

        self.__message_buffer += text

    def flush(self):

        if not self.__message_buffer:
            self.sendMessage('update')
            return

        self.__message_buffer = \
            self.__message_buffer.replace('\\', '\\\\').replace('"', '\\"')\
                .replace('\t', ' ').replace('-', '−')

        messages = iter(self.__message_buffer.split('\n'))

        self.__printMsg(next(messages))
        for message in messages:
            self.sendMessage('LF')
            if message:
                self.__printMsg(message)

        self.sendMessage('update')
        self.__message_buffer = ''

    def clear(self):
        self.sendMessage('clear')
        self.sendMessage('set-cursor-pos 0 0')
        self.__message_buffer = ''

    def __printMsg(self, text):

        if self.__current_message_mode is \
                SimpleConsoleOutputDevice.WriteMode.Normal:
            cmd = 'write'
        elif self.__current_message_mode is \
                SimpleConsoleOutputDevice.WriteMode.Static:
            cmd = 'static-write'
        elif self.__current_message_mode is \
                SimpleConsoleOutputDevice.WriteMode.Push:
            cmd = 'push-write'
        else:
            cmd = 'write'

        self.sendMessage(f'{cmd} "{text}"')

Device._DEVICE_TYPE_MAP['console-text-display'] = SimpleConsoleOutputDevice

class Ship:

    class ConsolePrinter:

        def __init__(self, ship):
            self.__ship = ship

        def write(self, text):
            self.__ship.writeConsole(text)

        def flush(self):
            pass

    class KeyboardReader:

        def __init__(self, ship):
            self.__ship = ship

        def readline(self):
            return self.__ship.readKeyboard() + '\n'

    def __init__(self):
        self.__device = Device()
        self.__sensor_devices = {}
        self.__interface_devices = {}
        self.__engine_devices = {}
        self.__console_printer = Ship.ConsolePrinter(self)
        self.__keyboard_reader = Ship.KeyboardReader(self)

        self.__find_devices(self.__device)

        console = self.__interface_devices.get('console-text-display')
        if console:
            self.__main_console = console[0]
        else:
            self.__main_console = None

        keyboard = self.__interface_devices.get('keyboard')
        if keyboard:
            self.__main_keyboard = keyboard[0]
        else:
            self.__main_keyboard = None

    def run(self, seconds):
        start_time = time.time()

        if self.__main_console is not None:
            self.__main_console.flush()

        time.sleep(max(seconds - (time.time() - start_time), 0))

    def listInterfaceDevices(self, device_type=None):
        if device_type is None:
            return self.__interface_devices

        return self.__interface_devices.get(device_type, ())

    def listEngines(self, engine_type=None):
        if engine_type is None:
            return self.__engine_devices

        return self.__engine_devices.get(engine_type, ())

    def listSensors(self, sensor_type=None):
        if sensor_type is None:
            return self.__sensor_devices

        return self.__sensor_devices.get(sensor_type, ())

    @property
    def console_printer(self):
        return self.__console_printer

    @property
    def keyboard_reader(self):
        return self.__keyboard_reader

    def writeConsole(self, text):
        self.__main_console.write(text)

    def readKeyboard(self):
        return self.__main_keyboard.readline()

    def displayPrint(self, message):

        text_displays = self.__interface_devices.get('text-display')
        if text_displays:
            text_displays[0].sendMessage(f'set-text "{message}"')

    @property
    def position(self):
        if not self.__sensor_devices:
            return None

        position_devices = self.__sensor_devices.get('position-sensor')

        if not position_devices:
            return None

        device = position_devices[0]

        return (device.read('x'), device.read('y'))

    @property
    def angle(self):
        if not self.__sensor_devices:
            return None

        position_devices = self.__sensor_devices.get('angle-sensor')

        if not position_devices:
            return None

        device = position_devices[0]

        return float(device.read())

    @property
    def speed(self):
        if not self.__sensor_devices:
            return None

        position_devices = self.__sensor_devices.get('speed-sensor')

        if not position_devices:
            return None

        device = position_devices[0]

        return device.read()

    @property
    def angular_speed(self):
        if not self.__sensor_devices:
            return None

        position_devices = self.__sensor_devices.get('ang-speed-sensor')

        if not position_devices:
            return None

        device = position_devices[0]

        return device.read()

    @property
    def device(self):
        return self.__device

    def __find_devices(self, device):

        if isinstance(device, Sensor):

            devices = self.__sensor_devices.get(device.type_)
            if devices is None:
                self.__sensor_devices[device.type_] = [device]
            else:
                devices.append(device)
            return

        if device.type_ in ('text-display', 'console-text-display',
                            'keyboard'):
            devices = self.__interface_devices.get(device.type_)
            if devices is None:
                self.__interface_devices[device.type_] = [device]
            else:
                devices.append(device)
            return

        if isinstance(device, Engine):
            devices = self.__engine_devices.get(device.type_)
            if devices is None:
                self.__engine_devices[device.type_] = [device]
            else:
                devices.append(device)
            return

        children = device.children
        if children is None:
            return

        for child in device.children:
            self.__find_devices(child)

def send(message):

    __device_comm_write.write(message)
    __device_comm_write.write('\n')
    __device_comm_write.flush()

    return __device_comm_read.readline()[:-1]

def debug(*args, **kwargs):
    print(*args, **kwargs, file=sys.stderr)
    sys.stderr.flush()

ship = Ship()

sys.stdout = ship.console_printer
sys.stdin = ship.keyboard_reader


from datetime import datetime, timedelta
import logging
import requests
from sqlalchemy import BigInteger, Column, DateTime, ForeignKey, SmallInteger, String, text, Text, desc, and_, exists, \
    or_
from sqlalchemy.orm import relationship
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import create_engine
import os
from sqlalchemy.orm import sessionmaker
from common.constant import intro_position
from common.constant.session_status import SessionStatus
from common.constant.user_status import UserStatus

Base = declarative_base()

try:
    db_url = os.environ.get('DATABASE_URL')
    engine = create_engine(db_url, echo=False)
    Base.metadata.create_all(engine)
    DBSession = sessionmaker(bind=engine)
except:
    logging.exception('')

class SessionContext(object):
    def __init__(self, session):
        self.session = session

    def __enter__(self):
        return self.session

    def __exit__(self, exc_type, exc_value, traceback):
        try:
            self.session.flush()
            self.session.commit()
            self.session.close()
        except:
            logging.exception('')
            self.session.close()

class SessionFactory(object):
    @staticmethod
    def create():
        try:
            return SessionContext(DBSession())
        except:
            logging.exception('')

class User(Base):
    __tablename__ = 'users'

    id = Column(BigInteger, primary_key=True)
    sender_id = Column(BigInteger, nullable=False, unique=True)
    first_name = Column(String(255), nullable=False)
    middle_name = Column(String(255), server_default=text("NULL::character varying"))
    last_name = Column(String(255), server_default=text("NULL::character varying"))
    nickname = Column(String(255), server_default=text("NULL::character varying"))
    status = Column(SmallInteger, nullable=False, server_default=text("0"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))

    @classmethod
    def find_by_id(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                user = db_session \
                    .query(
                    cls.sender_id,
                    cls.first_name,
                    cls.status
                ) \
                    .filter(cls.id == user_id).all()[0]

                return user
            except:
                logging.exception('')

    @classmethod
    def find_by_sender_id(cls, sender_id):
        with SessionFactory.create() as db_session:
            try:
                user_id = db_session \
                    .query(cls.id) \
                    .filter(cls.sender_id == sender_id).limit(1).all()[0][0]

                return user_id
            except:
                logging.exception('')

    @classmethod
    def register_user(cls, sender_id):
        with SessionFactory.create() as db_session:
            try:
                user_id = db_session \
                    .query(cls.id) \
                    .filter(cls.sender_id == sender_id).limit(1).all()

                if not user_id:
                    user_profile = cls.__get_user_profile(sender_id)

                    new_user = cls()
                    new_user.sender_id = sender_id
                    new_user.first_name = user_profile['first_name']
                    new_user.last_name = user_profile['last_name']

                    db_session.add(new_user)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def __get_user_profile(cls, fbid):
        try:
            url = "https://graph.facebook.com/v2.6/" + str(
                fbid) + "?fields=first_name,last_name,profile_pic,locale,timezone,gender&access_token=" + str(
                os.environ.get("PAGE_ACCESS_TOKEN", None))

            user_profile = requests.get(url).json()

            if user_profile is None:
                user_profile = cls.__get_unknown_profile_information()
            else:
                if 'error' in user_profile.keys():
                    user_profile = cls.__get_unknown_profile_information()

            return user_profile
        except:
            logging.exception('')

    @classmethod
    def __get_unknown_profile_information(cls):
        return {'first_name': 'UNKNOWN_FIRST_NAME', 'last_name': None}

    @classmethod
    def update_status(cls, status, user_id):
        with SessionFactory.create() as db_session:
            try:
                user_to_update = db_session.query(cls).filter(cls.id == user_id).all()[0]
                user_to_update.status = status
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def find_possible_remind_users(cls):
        with SessionFactory.create() as db_session:
            try:
                filter_1 = cls.created_at > datetime(2018, 5, 9)

                filter_2 = ~cls.status.in_(
                    [UserStatus.GET_STARTED.value, UserStatus.INTRO.value, UserStatus.SUICIDE_ILLNESS_INTRO.value,
                     UserStatus.SUICIDE_IN_SESSION.value])

                possible_remind_users = db_session.query(cls.id) \
                    .filter(and_(filter_1, filter_2)).order_by(cls.id).all()

                possible_remind_users = [i[0] for i in possible_remind_users]

                return possible_remind_users
            except:
                logging.exception('')

    @classmethod
    def find_unfinished_message_users(cls):
        with SessionFactory.create() as db_session:
            try:
                interval = 20

                unfinished_message_users = db_session \
                    .query(cls.id) \
                    .filter(
                    exists().where(
                        and_(
                            Message.user_id == cls.id,
                            Message.created_at > datetime.utcnow() - timedelta(seconds=interval)
                        )
                    )
                ).all()

                unfinished_message_users = [i[0] for i in unfinished_message_users]

                return unfinished_message_users
            except:
                logging.exception('')
                return []

    @classmethod
    def find_intro_users(cls, unread_messages):
        with SessionFactory.create() as db_session:
            try:
                intro_users = db_session \
                    .query(cls.id) \
                    .filter(
                    and_(
                        cls.id.in_([i[0] for i in unread_messages]),
                        cls.status.in_([0, 1, 7])
                    )
                ).all()

                if intro_users:
                    intro_users = [i[0] for i in intro_users]

                return intro_users
            except:
                logging.exception('')
                return []

    @classmethod
    def find_inactivated_user_ids(cls):
        with SessionFactory.create() as db_session:
            try:
                inactivated_sessions = db_session \
                    .query(Session.user_id) \
                    .filter(
                    and_(
                        Session.status == SessionStatus.active.value,
                        Session.finish_at < datetime.utcnow() - timedelta(minutes=5),
                        Session.created_at > datetime(year=2018, month=6, day=22)
                    )
                ).distinct(Session.user_id).all()

                if inactivated_sessions:
                    inactivated_user_ids = [i[0] for i in inactivated_sessions]
                else:
                    inactivated_user_ids = []

                print("\nINACTIVATED_USER\n{}".format(inactivated_user_ids))

                return inactivated_user_ids
            except:
                logging.exception('')
                return []

    @classmethod
    def find_sender_id_by_id(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                sender_id = db_session \
                    .query(cls.sender_id) \
                    .filter(cls.id == user_id).limit(1).all()[0][0]

                return sender_id
            except:
                logging.exception('')

    @classmethod
    def find_users_in_worker(cls, finished_messages):
        with SessionFactory.create() as db_session:
            try:
                finished_message_users = [i[0] for i in finished_messages]

                users_in_worker = db_session \
                    .query(cls.id) \
                    .filter(
                    and_(
                        cls.id.in_(finished_message_users),
                        exists().where(
                            and_(
                                Message.user_id == cls.id,
                                Message.read_flag == 1,
                                Message.status == 0,
                                Message.created_at > datetime.utcnow() - timedelta(minutes=3)
                            )
                        )
                    )
                ).all()

                users_in_worker = [i[0] for i in users_in_worker]

                if not users_in_worker:
                    return []
                else:
                    return users_in_worker
            except:
                logging.exception('')
                return []

    @classmethod
    def find_first_name_by_id(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                first_name = db_session \
                    .query(cls.first_name) \
                    .filter(cls.id == user_id).all()[0][0]

                return first_name
            except:
                logging.exception('')

class MessageCluster(Base):
    __tablename__ = 'message_clusters'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'), nullable=False)
    length = Column(String(255), server_default=text("NULL::character varying"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))

    user = relationship('User')

    @classmethod
    def register_new_message_cluster(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                new_message_cluster = cls()
                new_message_cluster.user_id = user_id

                db_session.add(new_message_cluster)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def find_id_by_user_id(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                cluster_id = db_session \
                    .query(cls.id) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).limit(1).all()

                cluster_id = cluster_id[0][0]

                return cluster_id
            except:
                logging.exception('')

    @classmethod
    def find_last_2_message_cluster_time(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                last_two_message_cluster_time = db_session \
                    .query(
                    cls.created_at
                ) \
                    .filter(
                    cls.user_id == user_id
                ).order_by(desc(cls.created_at)).limit(2).all()

                last_two_message_cluster_time = [i[0] for i in last_two_message_cluster_time]

                return last_two_message_cluster_time
            except:
                logging.exception('')
                return []

class Reaction(Base):
    __tablename__ = 'reactions'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'))
    reaction_id = Column(SmallInteger, nullable=False, server_default=text("0"))
    reaction_type = Column(String(255), nullable=False, server_default=text("0"))
    status = Column(SmallInteger, nullable=False, server_default=text("0"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))

    user = relationship('User')

    @classmethod
    def enable_old_reactions(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                reaction_to_update = db_session \
                    .query(cls) \
                    .filter(
                    and_(
                        cls.user_id == user_id,
                        cls.created_at < datetime.utcnow() - timedelta(hours=8)
                    )
                ).all()

                for reaction in reaction_to_update:
                    reaction.status = 1
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def enable_reaction_number(cls, user_id, reaction_type, used_list=None):
        with SessionFactory.create() as db_session:
            try:
                if used_list is None:
                    reaction_to_update = db_session \
                        .query(cls) \
                        .filter(
                        and_(
                            cls.user_id == user_id,
                            cls.reaction_type == reaction_type,
                            cls.status == 0
                        )
                    ).all()
                else:
                    reaction_to_update = db_session \
                        .query(cls) \
                        .filter(
                        and_(
                            cls.user_id == user_id,
                            cls.reaction_id.in_(used_list),
                            cls.status == 0
                        )
                    ).all()

                for reaction in reaction_to_update:
                    reaction.status = 1
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def disable_reaction_number(cls, user_id, reaction_id, reaction_type):
        with SessionFactory.create() as db_session:
            try:
                new_reaction_number = cls()

                new_reaction_number.user_id = user_id
                new_reaction_number.reaction_id = str(reaction_id)
                new_reaction_number.reaction_type = reaction_type

                db_session.add(new_reaction_number)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def find_used_reaction_number(cls, user_id, reaction_type):
        with SessionFactory.create() as db_session:
            try:
                cls.enable_old_reactions(user_id)

                used_reaction_number = db_session \
                    .query(cls.reaction_id) \
                    .filter(
                    and_(
                        cls.user_id == user_id,
                        cls.reaction_type == reaction_type,
                        cls.status == 0
                    )
                ).all()

                used_reaction_number = [i[0] for i in used_reaction_number]

                return used_reaction_number
            except:
                logging.exception('')
                return []

class Session(Base):
    __tablename__ = 'sessions'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'), nullable=False)
    status = Column(SmallInteger, nullable=False, server_default=text("0"))
    mood_start = Column(String(255), server_default=text("NULL::character varying"))
    mood_end = Column(String(255), server_default=text("NULL::character varying"))
    comment_end = Column(Text, server_default=text("NULL::character varying"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))
    finish_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))

    user = relationship('User')

    @classmethod
    def create_prepared_session(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                new_session = cls()
                new_session.user_id = user_id
                new_session.status = SessionStatus.prepared.value

                db_session.add(new_session)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def find_latest_id_by_user_id(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                user_session = db_session \
                    .query(cls) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).limit(1).all()

                if user_session:
                    return user_session[0].id
                else:
                    return None
            except:
                logging.exception('')

    @classmethod
    def find_latest_session_finish_at(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                latest_session_time = db_session \
                    .query(cls.finish_at) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).limit(1).all()[0][0]

                if latest_session_time:
                    return latest_session_time
                else:
                    return None
            except:
                logging.exception('')

    @classmethod
    def find_latest_session_data(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                user_session = db_session \
                    .query(cls.id, cls.user_id, cls.status, cls.finish_at) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).limit(1).all()

                if user_session:
                    user_session = [i._asdict() for i in user_session][0]
                    return user_session
                else:
                    return None
            except:
                logging.exception('')

    @classmethod
    def update_status(cls, session_id, status):
        with SessionFactory.create() as db_session:
            try:
                session_to_update = db_session \
                    .query(cls) \
                    .filter(cls.id == session_id).limit(1).all()[0]

                session_to_update.status = status
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def update_finish_at(cls, session_id, finish_at):
        with SessionFactory.create() as db_session:
            try:
                session_to_update = db_session \
                    .query(cls) \
                    .filter(cls.id == session_id).limit(1).all()[0]

                session_to_update.finish_at = finish_at
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def update_latest_session_status(cls, session_status: SessionStatus, user_id):
        with SessionFactory.create() as db_session:
            try:
                session_to_update = db_session \
                    .query(cls) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).limit(1).all()[0]

                if session_to_update:
                    session_to_update.status = session_status
            except:
                logging.exception('')

    @classmethod
    def save_comment_end(cls, session_id, message):
        with SessionFactory.create() as db_session:
            try:
                comment = ''.join([i['text'] for i in message.message_dicts])

                session_to_update = db_session \
                    .query(cls) \
                    .filter(
                    cls.id == session_id
                ).all()[0]

                print('\n[SAVING COMMENT]')

                session_to_update.comment_end = comment
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def save_mood_end(cls, user_id, mood: str):
        with SessionFactory.create() as db_session:
            try:
                session_to_update = db_session \
                    .query(cls) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).limit(1).all()[0]

                print('\n[SAVING MOOD]')

                session_to_update.mood_end = mood
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def update_existing_sessions(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                existing_sessions = db_session \
                    .query(cls) \
                    .filter(cls.user_id == user_id).all()

                for sess in existing_sessions:
                    sess.finish_at = datetime.utcnow() - timedelta(hours=9)
                    sess.status = SessionStatus.ended.value
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def create_new_session(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                new_session = cls()
                new_session.user_id = user_id
                new_session.finish_at = datetime.utcnow() + timedelta(minutes=30)

                db_session.add(new_session)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def admin_create_asking_comment_session(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                new_session = cls()
                new_session.user_id = user_id
                new_session.status = SessionStatus.asking_comment.value
                new_session.finish_at = datetime.utcnow() - timedelta(minutes=35)

                db_session.add(new_session)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def admin_update_status_asking_comment(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                sessions = db_session \
                    .query(cls) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).all()

                for idx, i in enumerate(sessions):
                    i.finish_at = datetime.utcnow() - timedelta(minutes=35)
                    if idx == 0:
                        i.status = SessionStatus.asking_comment.value
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def fetch_session_started_time(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                session_started_time = db_session \
                    .query(
                    Session.created_at
                ).filter(
                    Session.user_id == user_id,
                ).order_by(Session.created_at.desc()).all()[0][0]

                return session_started_time
            except:
                logging.exception('')
                return ""

    @staticmethod
    def fetch_session_ended_time(user_id):
        with SessionFactory.create() as db_session:
            try:
                session_ended_time = db_session \
                    .query(
                    Session.finish_at
                ).filter(
                    Session.user_id == user_id,
                ).order_by(Session.created_at.desc()).all()[0][0]

                return session_ended_time

            except:
                logging.exception('')
                return ""

class Message(Base):
    __tablename__ = 'messages'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'), nullable=False)
    session_id = Column(ForeignKey('sessions.id'), nullable=True)
    cluster_id = Column(ForeignKey('message_clusters.id'))
    message = Column(Text, nullable=False)
    status = Column(SmallInteger, nullable=False, server_default=text("0"))
    read_flag = Column(SmallInteger, nullable=False, server_default=text("0"))
    payload = Column(String(255), server_default=text("NULL::character varying"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))

    cluster = relationship('MessageCluster')
    user = relationship('User')
    db_session = relationship('Session')

    @classmethod
    def find_unread_messages(cls):
        with SessionFactory.create() as db_session:
            try:
                unread_messages = db_session \
                    .query(
                    cls.user_id,
                    cls.message,
                    cls.id
                ) \
                    .filter(
                    and_(
                        cls.read_flag == 0,
                        cls.status == 0
                    )
                ) \
                    .order_by(cls.user_id).all()

                return unread_messages
            except:
                logging.exception('')
                return []

    @classmethod
    def find_cluster_id(cls, message_dict):
        with SessionFactory.create() as db_session:
            try:
                message_id_list = [i['id'] for i in message_dict]

                cluster_id = db_session \
                    .query(
                    cls.cluster_id
                ) \
                    .filter(
                    cls.id.in_(message_id_list)
                ).limit(1).all()

                if cluster_id:
                    return cluster_id[0][0]
                else:
                    return None
            except:
                logging.exception('')

    @classmethod
    def find_latest_message_time(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                latest_message_time = db_session \
                    .query(
                    cls.created_at
                ) \
                    .filter(
                    cls.user_id == user_id
                ).order_by(desc(cls.id)).limit(1).all()

                if not latest_message_time:
                    return None
                else:
                    latest_message_time = latest_message_time[0][0]

                return latest_message_time
            except:
                logging.exception('')

    @classmethod
    def tag_msg_by_session_id(cls, session_id, message_id_list):
        with SessionFactory.create() as db_session:
            try:
                messages_to_update = db_session \
                    .query(cls) \
                    .filter(
                    cls.id.in_(message_id_list)
                ).all()

                for m in messages_to_update:
                    m.session_id = session_id
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def change_current_and_new_message_status(cls, message_id_list, user_id):
        with SessionFactory.create() as db_session:
            try:
                message_to_update = db_session \
                    .query(cls) \
                    .filter(
                    and_(
                        cls.id >= message_id_list[0],
                        cls.user_id == user_id
                    )
                ).all()

                for m in message_to_update:
                    m.read_flag = 0
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def has_user_sent_new_message(cls, user_id, message_id_list):
        with SessionFactory.create() as db_session:
            try:
                new_message = db_session \
                    .query(cls) \
                    .filter(
                    and_(
                        cls.user_id == user_id,
                        cls.id > max(message_id_list)
                    )
                ).all()

                if new_message:
                    return True
                else:
                    return False
            except:
                logging.exception('')
                return False

    @classmethod
    def find_second_last_message_time(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                second_last_message_time = db_session \
                    .query(
                    cls.created_at
                ) \
                    .filter(
                    and_(
                        cls.user_id == user_id,
                        cls.read_flag == 1
                    )
                ).order_by(desc(cls.id)).limit(2).all()

                return second_last_message_time[1][0]
            except:
                logging.exception('')
                return []

    @classmethod
    def update_read_flag(cls, status, message_id):
        with SessionFactory.create() as db_session:
            try:
                message_to_update = db_session \
                    .query(cls) \
                    .filter(cls.id == message_id).first()

                message_to_update.read_flag = status
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def handle_failed_messages(cls):
        # return a list like [[user_id, message_text, message_id], [1, 'hoge', 34]]
        with SessionFactory.create() as db_session:
            try:
                failed_messages = db_session \
                    .query(cls) \
                    .filter(
                    and_(
                        cls.status == 0,
                        cls.read_flag == 1,
                        cls.created_at < datetime.utcnow() - timedelta(minutes=10)
                    )
                ).order_by(cls.id).all()

                if not failed_messages:
                    return

                failed_message_users = list(set([i.id for i in failed_messages]))
                messages_to_unread = []
                messages_to_read = []
                for i in failed_message_users:
                    newest_successful_message_id = db_session \
                        .query(cls.id) \
                        .filter(
                        and_(
                            cls.user_id == i,
                            cls.status == 1
                        )
                    ).order_by(desc(cls.id)).limit(1).all()

                    this_user_messages = [i for i in failed_messages if i.id == i]
                    if any(i.id < newest_successful_message_id for i in this_user_messages):
                        messages_to_read.extend(this_user_messages)
                    else:
                        messages_to_unread.extend(this_user_messages)

                for i in messages_to_unread:
                    i.read_flag = 0

                for i in messages_to_read:
                    i.status = 1
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def change_message_status(cls, message_id_list):
        with SessionFactory.create() as db_session:
            try:
                message_to_update = db_session \
                    .query(
                    cls
                ) \
                    .filter(
                    and_(
                        cls.id.in_(message_id_list),
                        cls.status == 0
                    )
                ).all()

                for message in message_to_update:
                    message.status = 1
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def save_message(cls, user_id, created_at, message, payload=''):
        with SessionFactory.create() as db_session:
            try:
                new_message = cls()
                new_message.user_id = user_id
                new_message.message = message
                new_message.created_at = created_at
                new_message.payload = payload

                db_session.add(new_message)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def tag_messages_with_cluster_id(cls, message_id_list, cluster_id):
        with SessionFactory.create() as db_session:
            try:
                message_to_update = db_session \
                    .query(cls) \
                    .filter(
                    cls.id.in_(message_id_list)
                ).all()

                for message in message_to_update:
                    message.cluster_id = cluster_id
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def fetch_previous_msg(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                latest_cluster_id = db_session \
                    .query(
                    Message.cluster_id
                ) \
                    .filter(
                    Message.user_id == user_id
                ) \
                    .order_by(
                    desc(Message.created_at)).limit(1).all()[0][0]

                if latest_cluster_id is None:
                    return ""

                previous_msg_list = db_session \
                    .query(
                    Message.message
                ) \
                    .filter(
                    Message.user_id == user_id,
                    Message.cluster_id == latest_cluster_id - 1
                ) \
                    .all()[0]

                if len(previous_msg_list) == 0 or previous_msg_list is None:
                    return ""

                else:
                    return previous_msg_list[0]

            except:
                logging.exception('')
                return ""

    @classmethod
    def count_total_msg_in_session(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                session_started_time = Session.fetch_session_started_time(user_id)
                session_ended_time = Session.fetch_session_ended_time(user_id)
                total_number_of_msg_in_session = db_session \
                    .query(
                    Message
                ) \
                    .filter(
                    Message.user_id == user_id,
                    Message.created_at > session_started_time,
                    Message.created_at < session_ended_time
                ).count()

                return total_number_of_msg_in_session

            except:
                logging.exception('')
                return ""

class Response(Base):
    __tablename__ = 'responses'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'))
    cluster_id = Column(ForeignKey('message_clusters.id'))
    type = Column(String(255), nullable=False)
    response = Column(Text, nullable=False)
    sent_at = Column(DateTime, nullable=False)
    sent_flag = Column(SmallInteger, nullable=False, server_default=text("0"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))
    cluster = relationship('MessageCluster')
    user = relationship('User')

    @classmethod
    def save_response_data(cls, user_id, cluster_id, response, message_types: list, delay=0, has_sent=False):
        with SessionFactory.create() as db_session:
            try:
                new_response = cls()
                new_response.user_id = user_id
                new_response.cluster_id = cluster_id
                new_response.response = response
                new_response.type = ' '.join(message_types)

                if has_sent:
                    new_response.sent_at = datetime.utcnow()
                    new_response.sent_flag = 1
                else:
                    new_response.sent_at = datetime.utcnow() + timedelta(seconds=delay)
                    new_response.sent_flag = 0

                db_session.add(new_response)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def exists_future_message(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                response_id = db_session \
                    .query(cls.id) \
                    .filter(
                    and_(
                        cls.user_id == user_id,
                        cls.sent_flag == 0
                    )
                ).order_by(cls.sent_at).limit(1).all()

                if response_id:
                    return True
                else:
                    return False
            except:
                logging.exception('')
                return False

    @classmethod
    def find_unsent_message_data(cls):
        with SessionFactory.create() as db_session:
            try:
                unsent_messages = db_session \
                    .query(
                    cls.id,
                    cls.user_id,
                    cls.response
                ) \
                    .filter(
                    and_(
                        cls.sent_at < datetime.utcnow(),
                        cls.sent_flag == 0
                    )
                ).order_by(cls.id).all()

                return unsent_messages
            except:
                logging.exception('')
                return []

    @classmethod
    def update_response_sent_flag(cls, unsent_messages):
        with SessionFactory.create() as db_session:
            try:
                response_id_list = [i[0] for i in unsent_messages]

                response_to_update = db_session \
                    .query(cls) \
                    .filter(
                    and_(
                        cls.sent_flag == 0,
                        or_(
                            cls.id.in_(response_id_list),
                            cls.sent_at < datetime.utcnow() - timedelta(minutes=10),
                            cls.sent_at > datetime.utcnow() + timedelta(minutes=10)
                        )
                    )
                ).all()

                for response in response_to_update:
                    response.sent_flag = 1
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def find_past_3_response_types(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                last_3_cluster_ids = db_session \
                    .query(MessageCluster.id) \
                    .filter(MessageCluster.user_id == user_id).order_by(desc(MessageCluster.id)).limit(3).all()

                last_3_cluster_ids = [i[0] for i in last_3_cluster_ids]

                past_response_types = db_session \
                    .query(
                    cls.cluster_id,
                    cls.type
                ) \
                    .filter(cls.cluster_id.in_(last_3_cluster_ids)).order_by(cls.cluster_id).all()

                if past_response_types:
                    response_types = []
                    for cluster_id in last_3_cluster_ids:
                        r_types = [i[1].split(' ') for i in past_response_types if i[0] == cluster_id]

                        if r_types:
                            response_types += r_types[0]

                    return response_types
                else:
                    return []
            except:
                logging.exception('')
                return []

    @classmethod
    def fetch_last_response_type_list(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                last_response_types = db_session \
                    .query(
                    Response.type
                ) \
                    .filter(
                    Response.user_id == user_id
                ).order_by(desc(Response.created_at)).limit(3).all()[0][0].split()

                return last_response_types
            except:
                logging.exception('')
                return []

    @classmethod
    def fetch_all_response_type_during(cls, session_started_time, session_ended_time, user_id):
        with SessionFactory.create() as db_session:
            try:
                past_response_types = db_session \
                    .query(
                    Response.cluster_id,
                    Response.type
                ) \
                    .filter(
                    Response.user_id == user_id,
                    Response.created_at > session_started_time,
                    Response.created_at < session_ended_time
                ).all()

                return past_response_types
            except:
                logging.exception('')
                return []

    @classmethod
    def fetch_previous_msg_response_type(cls, user_id):
        # TODO see the previous response type and return correct response type
        with SessionFactory.create() as db_session:
            try:
                previous__msg_response_type_list = db_session \
                    .query(
                    Response.type
                ) \
                    .filter(
                    Response.user_id == user_id
                ) \
                    .order_by(desc(Response.id)).limit(1).all()[0]

                if len(previous__msg_response_type_list) == 0 or previous__msg_response_type_list is None:
                    return []
                else:
                    return previous__msg_response_type_list[0].split()

            except:
                logging.exception('')
                return []

class Remind(Base):
    __tablename__ = 'reminds'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'))
    type = Column(SmallInteger, nullable=False, server_default=text("0"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))

    user = relationship('User')

    @classmethod
    def register_remind(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                new_remind = cls()
                new_remind.user_id = user_id

                db_session.add(new_remind)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def find_latest_remind_time(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                latest_remind_time = db_session \
                    .query(
                    cls.created_at
                ) \
                    .filter(
                    cls.user_id == user_id
                ).order_by(desc(cls.id)).limit(1).all()

                if latest_remind_time:
                    return latest_remind_time[0]
                else:
                    return None
            except:
                logging.exception('')

    @classmethod
    def find_reminds_by_created_at(cls, user_id, time):
        with SessionFactory.create() as db_session:
            try:
                # TODO can't pass remind class. find remind here
                remind = db_session \
                    .query(cls) \
                    .filter(
                    and_(
                        cls.user_id == user_id,
                        cls.created_at > time
                    )
                ).order_by(desc(cls.id)).all()

                return remind
            except:
                logging.exception('')

    @classmethod
    def register_remind_mood(cls, mood: str, user_id):
        # TODO can't pass remind class. find remind here
        with SessionFactory.create() as db_session:
            try:
                latest_remind = db_session \
                    .query(cls) \
                    .filter(cls.user_id == user_id).order_by(desc(cls.id)).limit(1).all()[0]

                latest_remind.type = mood
            except:
                db_session.rollback()
                logging.exception('')

class UsersFeeling(Base):
    __tablename__ = 'users_feelings'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'))
    keyword = Column(String(255), nullable=False)
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))

    user = relationship('User')

    @classmethod
    def save_feelings(cls, text_kw_df, user_id):
        with SessionFactory.create() as db_session:
            try:
                if text_kw_df is None:
                    return []

                negative_feeling_df = text_kw_df[text_kw_df.fact.isin([False]) & (text_kw_df.sscore < 0)]

                if negative_feeling_df.empty:
                    return

                for idx, row in negative_feeling_df.iterrows():
                    new_feeling = cls()
                    new_feeling.user_id = user_id
                    new_feeling.keyword = row.word

                    db_session.add(new_feeling)
            except:
                logging.exception('')

class IntroPosition(Base):
    __tablename__ = 'intro_positions'

    id = Column(BigInteger, primary_key=True)
    user_id = Column(ForeignKey('users.id'))
    position = Column(String(255), nullable=False, server_default=text("0"))
    created_at = Column(DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP"))
    finished_at = Column(DateTime, server_default=text("CURRENT_TIMESTAMP"))

    user = relationship('User')

    @classmethod
    def register_new_intro(cls, user):
        with SessionFactory.create() as db_session:
            try:
                new_intro_position = cls()
                new_intro_position.user_id = user.id
                new_intro_position.position = intro_position.IntroPosition.GREETING.value

                db_session.add(new_intro_position)
            except:
                db_session.rollback()
                logging.exception('')

    @classmethod
    def find_position_by_user_id(cls, user_id):
        with SessionFactory.create() as db_session:
            try:
                position = db_session \
                    .query(cls.position) \
                    .filter(cls.user_id == user_id).limit(1).all()

                position = position[0][0]

                return position
            except:
                logging.exception('')

    @classmethod
    def save_position(cls, position, user_id):
        with SessionFactory.create() as db_session:
            try:
                position_to_update = db_session \
                    .query(cls) \
                    .filter(cls.user_id == user_id).first()

                position_to_update.position = position
            except:
                db_session.rollback()
                logging.exception('')

# metadata.create_all()

def delete_user_from_table(user_id, table):
    with SessionFactory.create() as db_session:
        try:
            if table == User:
                db_session.query(table).filter(table.id == user_id).delete()
            else:
                db_session.query(table).filter(table.user_id == user_id).delete()
        except:
            db_session.rollback()
            logging.exception('')


import pytest

from zmq_tubes.matcher import TopicMatcher

class TestTopicMatcher:

    @pytest.mark.parametrize("sub,topic", [
        ("foo/bar", "foo/bar"),
        ("foo/+", "foo/bar"),
        ("foo/+/baz", "foo/bar/baz"),
        ("foo/+/#", "foo/bar/baz"),
        ("A/B/+/#", "A/B/B/C"),
        ("#", "foo/bar/baz"),
        ("#", "/foo/bar"),
        ("/#", "/foo/bar"),
        ("$SYS/bar", "$SYS/bar"),
    ])
    def test_matching(self, sub, topic):
        t = TopicMatcher()
        t.set_topic(sub, True)
        assert t.match(topic)

    @pytest.mark.parametrize("sub,topic", [
        ("test/6/#", "test/3"),
        ("foo/bar", "foo"),
        ("foo/+", "foo/bar/baz"),
        ("foo/+/baz", "foo/bar/bar"),
        ("foo/+/#", "fo2/bar/baz"),
        ("/#", "foo/bar"),
        ("#", "$SYS/bar"),
        ("$BOB/bar", "$SYS/bar"),
    ])
    def test_not_matching(self, sub, topic):
        t = TopicMatcher()
        t.set_topic(sub, True)
        assert t.match(topic) is None

    @pytest.mark.parametrize("filter,val", [
        ('#', [('g/p/+/i', 'i'), ('g/p/+/o', 'o'), ('g/s', 's')]),
        ('g/#', [('g/p/+/i', 'i'), ('g/p/+/o', 'o'), ('g/s', 's')]),
        ('g/s', [('g/s', 's')]),
        ('g/+/+/i', [('g/p/+/i', 'i')]),
        ('g/+/+/+', [('g/p/+/i', 'i'), ('g/p/+/o', 'o')]),
        ('g/+', [('g/s', 's')]),
    ])
    def test_filtering(self, filter, val):
        t = TopicMatcher()
        t.set_topic('g/p/+/i', 'i')
        t.set_topic('g/p/+/o', 'o')
        t.set_topic('g/s', 's')
        assert set(t.filter(filter)) == set(val)


from collections import defaultdict

import torch
import torch.nn as nn

import numpy as np

class SplitCrossEntropyLoss(nn.Module):
    r'''SplitCrossEntropyLoss calculates an approximate softmax'''
    def __init__(self, hidden_size, splits, verbose=False):
        # We assume splits is [0, split1, split2, N] where N >= |V|
        # For example, a vocab of 1000 words may have splits [0] + [100, 500] + [inf]
        super(SplitCrossEntropyLoss, self).__init__()
        self.hidden_size = hidden_size
        self.splits = [0] + splits + [100 * 1000000]
        self.nsplits = len(self.splits) - 1
        self.stats = defaultdict(list)
        self.verbose = verbose
        # Each of the splits that aren't in the head require a pretend token, we'll call them tombstones
        # The probability given to this tombstone is the probability of selecting an item from the represented split
        if self.nsplits > 1:
            self.tail_vectors = nn.Parameter(torch.zeros(self.nsplits - 1, hidden_size))
            self.tail_bias = nn.Parameter(torch.zeros(self.nsplits - 1))

    def logprob(self, weight, bias, hiddens, splits=None, softmaxed_head_res=None, verbose=False):
        # First we perform the first softmax on the head vocabulary and the tombstones
        if softmaxed_head_res is None:
            start, end = self.splits[0], self.splits[1]
            head_weight = None if end - start == 0 else weight[start:end]
            head_bias = None if end - start == 0 else bias[start:end]
            # We only add the tombstones if we have more than one split
            if self.nsplits > 1:
                head_weight = self.tail_vectors if head_weight is None else torch.cat([head_weight, self.tail_vectors])
                head_bias = self.tail_bias if head_bias is None else torch.cat([head_bias, self.tail_bias])

            # Perform the softmax calculation for the word vectors in the head for all splits
            # We need to guard against empty splits as torch.cat does not like random lists
            head_res = torch.nn.functional.linear(hiddens, head_weight, bias=head_bias)
            softmaxed_head_res = torch.nn.functional.log_softmax(head_res)

        if splits is None:
            splits = list(range(self.nsplits))

        results = []
        running_offset = 0
        for idx in splits:

            # For those targets in the head (idx == 0) we only need to return their loss
            if idx == 0:
                results.append(softmaxed_head_res[:, :-(self.nsplits - 1)])

            # If the target is in one of the splits, the probability is the p(tombstone) * p(word within tombstone)
            else:
                start, end = self.splits[idx], self.splits[idx + 1]
                tail_weight = weight[start:end]
                tail_bias = bias[start:end]

                # Calculate the softmax for the words in the tombstone
                tail_res = torch.nn.functional.linear(hiddens, tail_weight, bias=tail_bias)

                # Then we calculate p(tombstone) * p(word in tombstone)
                # Adding is equivalent to multiplication in log space
                head_entropy = (softmaxed_head_res[:, -idx]).contiguous()
                tail_entropy = torch.nn.functional.log_softmax(tail_res)
                results.append(head_entropy.view(-1, 1) + tail_entropy)

        if len(results) > 1:
            return torch.cat(results, dim=1)
        return results[0]

    def split_on_targets(self, hiddens, targets):
        # Split the targets into those in the head and in the tail
        split_targets = []
        split_hiddens = []

        # Determine to which split each element belongs (for each start split value, add 1 if equal or greater)
        # This method appears slower at least for WT-103 values for approx softmax
        #masks = [(targets >= self.splits[idx]).view(1, -1) for idx in range(1, self.nsplits)]
        #mask = torch.sum(torch.cat(masks, dim=0), dim=0)
        ###
        # This is equally fast for smaller splits as method below but scales linearly
        mask = None
        for idx in range(1, self.nsplits):
            partial_mask = targets >= self.splits[idx]
            mask = mask + partial_mask if mask is not None else partial_mask
        ###
        #masks = torch.stack([targets] * (self.nsplits - 1))
        #mask = torch.sum(masks >= self.split_starts, dim=0)
        for idx in range(self.nsplits):
            # If there are no splits, avoid costly masked select
            if self.nsplits == 1:
                split_targets, split_hiddens = [targets], [hiddens]
                continue
            # If all the words are covered by earlier targets, we have empties so later stages don't freak out
            if sum(len(t) for t in split_targets) == len(targets):
                split_targets.append([])
                split_hiddens.append([])
                continue
            # Are you in our split?
            tmp_mask = mask == idx
            split_targets.append(torch.masked_select(targets, tmp_mask))
            split_hiddens.append(hiddens.masked_select(tmp_mask.unsqueeze(1).expand_as(hiddens)).view(-1, hiddens.size(1)))
        return split_targets, split_hiddens

    def forward(self, weight, bias, hiddens, targets, verbose=False):
        if self.verbose or verbose:
            for idx in sorted(self.stats):
                print('{}: {}'.format(idx, int(np.mean(self.stats[idx]))), end=', ')
            print()

        total_loss = None
        if len(hiddens.size()) > 2: hiddens = hiddens.view(-1, hiddens.size(2))

        split_targets, split_hiddens = self.split_on_targets(hiddens, targets)

        # First we perform the first softmax on the head vocabulary and the tombstones
        start, end = self.splits[0], self.splits[1]
        head_weight = None if end - start == 0 else weight[start:end]
        head_bias = None if end - start == 0 else bias[start:end]

        # We only add the tombstones if we have more than one split
        if self.nsplits > 1:
            head_weight = self.tail_vectors if head_weight is None else torch.cat([head_weight, self.tail_vectors])
            head_bias = self.tail_bias if head_bias is None else torch.cat([head_bias, self.tail_bias])

        # Perform the softmax calculation for the word vectors in the head for all splits
        # We need to guard against empty splits as torch.cat does not like random lists
        combo = torch.cat([split_hiddens[i] for i in range(self.nsplits) if len(split_hiddens[i])])
        ###
        all_head_res = torch.nn.functional.linear(combo, head_weight, bias=head_bias)
        softmaxed_all_head_res = torch.nn.functional.log_softmax(all_head_res)
        if self.verbose or verbose:
            self.stats[0].append(combo.size()[0] * head_weight.size()[0])

        running_offset = 0
        for idx in range(self.nsplits):
            # If there are no targets for this split, continue
            if len(split_targets[idx]) == 0: continue

            # For those targets in the head (idx == 0) we only need to return their loss
            if idx == 0:
                softmaxed_head_res = softmaxed_all_head_res[running_offset:running_offset + len(split_hiddens[idx])]
                entropy = -torch.gather(softmaxed_head_res, dim=1, index=split_targets[idx].view(-1, 1))
            # If the target is in one of the splits, the probability is the p(tombstone) * p(word within tombstone)
            else:
                softmaxed_head_res = softmaxed_all_head_res[running_offset:running_offset + len(split_hiddens[idx])]

                if self.verbose or verbose:
                    start, end = self.splits[idx], self.splits[idx + 1]
                    tail_weight = weight[start:end]
                    self.stats[idx].append(split_hiddens[idx].size()[0] * tail_weight.size()[0])

                # Calculate the softmax for the words in the tombstone
                tail_res = self.logprob(weight, bias, split_hiddens[idx], splits=[idx], softmaxed_head_res=softmaxed_head_res)

                # Then we calculate p(tombstone) * p(word in tombstone)
                # Adding is equivalent to multiplication in log space
                head_entropy = softmaxed_head_res[:, -idx]
                # All indices are shifted - if the first split handles [0,...,499] then the 500th in the second split will be 0 indexed
                indices = (split_targets[idx] - self.splits[idx]).view(-1, 1)
                # Warning: if you don't squeeze, you get an N x 1 return, which acts oddly with broadcasting
                tail_entropy = torch.gather(torch.nn.functional.log_softmax(tail_res), dim=1, index=indices).squeeze()
                entropy = -(head_entropy + tail_entropy)
            ###
            running_offset += len(split_hiddens[idx])
            total_loss = entropy.float().sum() if total_loss is None else total_loss + entropy.float().sum()

        return (total_loss / len(targets)).type_as(weight)

if __name__ == '__main__':
    np.random.seed(42)
    torch.manual_seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(42)

    V = 8
    H = 10
    N = 100
    E = 10

    embed = torch.nn.Embedding(V, H)
    crit = SplitCrossEntropyLoss(hidden_size=H, splits=[V // 2])
    bias = torch.nn.Parameter(torch.ones(V))
    optimizer = torch.optim.SGD(list(embed.parameters()) + list(crit.parameters()), lr=1)

    for _ in range(E):
        prev = (torch.rand(N, 1) * 0.999 * V).int().long()
        x = (torch.rand(N, 1) * 0.999 * V).int().long()
        y = embed(prev).squeeze()
        c = crit(embed.weight, bias, y, x.view(N))
        print('Crit', c.exp().item())

        logprobs = crit.logprob(embed.weight, bias, y[:2]).exp()
        print(logprobs)
        print(logprobs.sum(dim=1))

        optimizer.zero_grad()
        c.backward()
        optimizer.step()


#!env python
from mechanize import Browser
from urllib.error import HTTPError, URLError
import re
import sys
import urllib.parse

base = "http://localhost:5000"

def click_all_links(brower, url, seen_links, depth=0):
    full_url = urllib.parse.urljoin(base, url)
    if re.match(r"^[mailto.*|#.*]", url):
        return None

    leader = "-" * (depth + 1)
    print("%s Opening %s" % (leader, full_url))
    try:
        br.open(full_url)
    except (URLError, HTTPError) as e:
        print(e)
        sys.exit(1)

    for link in br.links():
        if link.url in ["/login", "/logout", "/admin"]:
            continue

        # skip player links
        if re.match(r"\/p\/\w{8}",  link.url):
            continue

        if link.url not in seen_links:
            seen_links.append(link.url)
            new_depth = depth + 1
            click_all_links(br, link.url, seen_links, new_depth)
        else:
            continue

    return len(seen_links)

br = Browser()
length = click_all_links(br, "", [])
print("Scanned %s links without issue" % length)