import numpy as np
import os
import pytest
import random
import tensorflow as tf
import torch
import utils

__author__ = "Christopher Potts"
__version__ = "CS224u, Stanford, Spring 2019"

tf.enable_eager_execution()

@pytest.mark.parametrize("arg, expected", [
    [
        np.array([0.0, 0.25, 0.75]),
        np.array([0.22721977, 0.29175596, 0.48102426])
    ]
])
def test_softmax(arg, expected):
    result = utils.softmax(arg).round(8)
    expected = expected.round(8)
    assert np.array_equal(result, expected)

@pytest.mark.parametrize("arg, expected", [
    [-1, 0],
    [np.array([-1.0, 1.0]), np.array([0.0, 0.0])]
])
def test_d_tanh(arg, expected):
    assert np.array_equal(utils.d_tanh(arg), expected)

def test_randvec():
    x = utils.randvec(10)
    assert len(x) == 10

def test_randmatrix():
    X = utils.randmatrix(10, 20)
    assert X.shape == (10, 20)

def test_safe_macro_f1():
    y = [1, 1, 2, 2, 1]
    y_pred = [1, 2, 2, 1, 1]
    utils.safe_macro_f1(y, y_pred)

@pytest.mark.parametrize("arg, expected", [
    [
        np.array([[1.0, 0.0], [0.0, 1.0]]),
        np.array([[0.0, 0.0], [0.0, 0.0]])
    ]
])
def test_log_of_array_ignoring_zeros(arg, expected):
    result = utils.log_of_array_ignoring_zeros(arg)
    return np.array_equal(result, expected)

def test_glove2dict():
    src_filename = os.path.join("data", "glove.6B", "glove.6B.50d.txt")
    data = utils.glove2dict(src_filename)
    assert len(data) == 400000

@pytest.mark.parametrize("X, n_words, expected", [
    [
        [["a", "b", "c"], ["b", "c", "d"]],
        None,
        ["$UNK", "a", "b", "c", "d"]
    ],
    [
        [["a", "b", "c"], ["b", "c", "d"]],
        2,
        ["$UNK", "b", "c"]
    ],
    [
        [],
        2,
        ["$UNK"]
    ]
])
def test_get_vocab(X, n_words, expected):
    result = utils.get_vocab(X, n_words=n_words)
    assert result == expected

@pytest.mark.parametrize("set_value", [True, False])
def test_fix_random_seeds_system(set_value):
    utils.fix_random_seeds(seed=42, set_system=set_value)
    x = np.random.random()
    utils.fix_random_seeds(seed=42, set_system=set_value)
    y = np.random.random()
    assert (x == y) == set_value

@pytest.mark.parametrize("set_value", [True, False])
def test_fix_random_seeds_pytorch(set_value):
    utils.fix_random_seeds(seed=42, set_torch=set_value)
    x = torch.rand(1)
    utils.fix_random_seeds(seed=42, set_torch=set_value)
    y = torch.rand(1)
    assert (x == y) == set_value

@pytest.mark.parametrize("set_value", [True, False])
def test_fix_random_seeds_tensorflow(set_value):
    utils.fix_random_seeds(seed=42, set_tensorflow=set_value)
    x = tf.random.uniform([1]).numpy()
    utils.fix_random_seeds(seed=42, set_tensorflow=set_value)
    y = tf.random.uniform([1]).numpy()
    assert (x == y) == set_value


import argparse
import tensorflow as tf
import os

# Freeze graph. Replaces variables to constants. Removes training-only ops.
def freeze_graph(net, ckpt, tool, out_node, out_graph='frozen_graph.pb'):
    os.system('python ' + tool +
              ' --input_graph ' + net +
              ' --input_checkpoint ' + ckpt +
              ' --output_graph ' + out_graph +
              ' --output_node_names ' + out_node)

def optimize_for_inference(dtype, tool, in_node, out_node, out_graph='optimized_graph.pb',
                           frozen_graph='frozen_graph.pb'):
    os.system('python ' + tool +
              ' --input ' + frozen_graph +
              ' --output ' + out_graph +
              ' --frozen_graph True ' +
              ' --input_names ' + in_node +
              ' --output_names ' + out_node +
              ' --placeholder_type_enum ' + str(dtype.as_datatype_enum))

def fuse_constants(tool, in_graph, out_graph, in_node, out_node):
    os.system(tool + ' --in_graph=' + out_graph + \
                     ' --out_graph=' + out_graph + \
                     ' --inputs=' + in_node + \
                     ' --outputs=' + out_node + \
                     ' --transforms="fold_constants(ignore_errors=True) sort_by_execution_order"')

def prepare_for_dnn(net, ckpt, freeze_graph_tool, optimizer_tool, transform_graph_tool,
                    input_node_name, output_node_name, out_graph, dtype):
    freeze_graph(net, ckpt, freeze_graph_tool, output_node_name)
    optimize_for_inference(dtype, optimizer_tool, input_node_name, output_node_name, out_graph)
    fuse_constants(transform_graph_tool, out_graph, out_graph, input_node_name, output_node_name)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Script for preparing serialized '
                                                 'TensorFlow graph to import into DNN. '
                                                 'Modified graph still may be used in TensorFlow.')
    parser.add_argument('-frz', dest='freeze_graph_tool', required=True,
                        help='Path to freeze_graph.py tool')
    parser.add_argument('-opt', dest='optimizer_tool', required=True,
                        help='Path to optimize_for_inference.py tool')
    parser.add_argument('-tr', dest='transform_graph_tool', required=True,
                        help='Path to transform_graph tool')
    parser.add_argument('-net', dest='net', required=True,
                        help='Path serialized graph by tf.train.write_graph()')
    parser.add_argument('-ckpt', dest='ckpt', required=True,
                        help='Path saved checkpoint by saver.save()')
    parser.add_argument('-in_node', dest='input_name', required=True,
                        help='Input op name')
    parser.add_argument('-out_node', dest='output_name', required=True,
                        help='Output op name')
    parser.add_argument('-o', dest='output', required=True,
                        help='Output graph name')
    args = parser.parse_args()

    prepare_for_dnn(args.net, args.ckpt, args.freeze_graph_tool,
                    args.optimizer_tool, args.transform_graph_tool,
                    args.input_name, args.output_name, args.output)


from __future__ import print_function
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import numpy as _np
import multiprocessing as _multiprocessing
import scipy.sparse as _sparse

class BootstrapResults(object):
    def __init__(self, lower_bound, value, upper_bound):
        self.lower_bound = lower_bound
        self.upper_bound = upper_bound
        self.value = value
        if self.lower_bound > self.upper_bound:
            self.lower_bound, self.upper_bound = self.upper_bound, self.lower_bound

    def __str__(self):
        return '{1}    ({0}, {2})'.format(self.lower_bound, self.value,
                                          self.upper_bound)

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

    def _apply(self, other, func):
        return BootstrapResults(func(self.lower_bound, other),
                                func(self.value, other),
                                func(self.upper_bound, other))

    def __add__(self, other):
        return self._apply(float(other), lambda x, other: other + x)

    def __radd__(self, other):
        return self._apply(float(other), lambda x, other: other + x)

    def __sub__(self, other):
        return self._apply(float(other), lambda x, other: x - other)

    def __rsub__(self, other):
        return self._apply(float(other), lambda x, other: other - x)

    def __mul__(self, other):
        return self._apply(float(other), lambda x, other: x * other)

    def __rmul__(self, other):
        return self._apply(float(other), lambda x, other: x * other)

    def error_width(self):
        '''Returns: upper_bound - lower_bound'''
        return self.upper_bound - self.lower_bound

    def error_fraction(self):
        '''Returns the error_width / value'''
        if self.value == 0:
            return _np.inf
        else:
            return self.error_width() / self.value

    def is_significant(self):
        return _np.sign(self.upper_bound) == _np.sign(self.lower_bound)

    def get_result(self):
        '''Returns:
            -1 if statistically significantly negative
            +1 if statistically significantly positive
            0 otherwise
        '''
        return int(self.is_significant()) * _np.sign(self.value)

def _get_confidence_interval(bootstrap_dist, stat_val, alpha, is_pivotal):
    '''Get the bootstrap confidence interval for a given distribution.
    Args:
        bootstrap_distribution: numpy array of bootstrap results from
            bootstrap_distribution() or bootstrap_ab_distribution()
        stat_val: The overall statistic that this method is attempting to
            calculate error bars for.
        alpha: The alpha value for the confidence intervals.
        is_pivotal: if true, use the pivotal method. if false, use the
            percentile method.
    '''
    if is_pivotal:
        low = 2 * stat_val - _np.percentile(bootstrap_dist, 100 * (1 - alpha / 2.))
        val = stat_val
        high = 2 * stat_val - _np.percentile(bootstrap_dist, 100 * (alpha / 2.))
    else:
        low = _np.percentile(bootstrap_dist, 100 * (alpha / 2.))
        val = _np.percentile(bootstrap_dist, 50)
        high = _np.percentile(bootstrap_dist, 100 * (1 - alpha / 2.))

    return BootstrapResults(low, val, high)

def _needs_sparse_unification(values_lists):
    non_zeros = values_lists[0] != 0

    for v in values_lists:
        v_nz = v != 0
        non_zeros = (non_zeros + v_nz) > 0

    non_zero_size = non_zeros.sum()

    for v in values_lists:
        if non_zero_size != v.data.shape[0]:
            return True

    return False

def _validate_arrays(values_lists):
    t = values_lists[0]
    t_type = type(t)
    if not isinstance(t, _sparse.csr_matrix) and not isinstance(t, _np.ndarray):
        raise ValueError(('The arrays must either be of type '
                          'scipy.sparse.csr_matrix or numpy.array'))

    for _, values in enumerate(values_lists[1:]):
        if not isinstance(values, t_type):
            raise ValueError('The arrays must all be of the same type')

        if t.shape != values.shape:
            raise ValueError('The arrays must all be of the same shape')

        if isinstance(t, _sparse.csr_matrix):
            if values.shape[0] > 1:
                raise ValueError(('The sparse matrix must have shape 1 row X N'
                                  ' columns'))

    if isinstance(t, _sparse.csr_matrix):
        if _needs_sparse_unification(values_lists):
            raise ValueError(('The non-zero entries in the sparse arrays'
                              ' must be aligned: see '
                              'bootstrapped.unify_sparse_vectors function'))

def _generate_distributions(values_lists, num_iterations):

    if isinstance(values_lists[0], _sparse.csr_matrix):
        # in the sparse case we dont actually need to bootstrap
        # the full sparse array since most values are 0
        # instead for each bootstrap iteration we:
        #    1. generate B number of non-zero entries to sample from the
        #          binomial distribution
        #    2. resample with replacement the non-zero entries from values
        #          B times
        #    3. create a new sparse array with the B resamples, zero otherwise
        results = [[] for _ in range(len(values_lists))]

        pop_size = values_lists[0].shape[1]
        non_sparse_size = values_lists[0].data.shape[0]

        p = non_sparse_size * 1.0 / pop_size

        for _ in range(num_iterations):
            ids = _np.random.choice(
                non_sparse_size,
                _np.random.binomial(pop_size, p),
                replace=True,
            )

            for arr, values in zip(results, values_lists):
                data = values.data
                d = _sparse.csr_matrix(
                    (
                        data[ids],
                        (_np.zeros_like(ids), _np.arange(len(ids)))
                    ),
                    shape=(1, pop_size),
                )

                arr.append(d)
        return [_sparse.vstack(r) for r in results]

    else:
        values_shape = values_lists[0].shape[0]
        ids = _np.random.choice(
            values_shape,
            (num_iterations, values_shape),
            replace=True,
        )

        results = [values[ids] for values in values_lists]
        return results

def _bootstrap_sim(values_lists, stat_func_lists, num_iterations,
                   iteration_batch_size, seed):
    '''Returns simulated bootstrap distribution.
    See bootstrap() funciton for arg descriptions.
    '''

    if seed is not None:
        _np.random.seed(seed)

    num_iterations = int(num_iterations)
    iteration_batch_size = int(iteration_batch_size)

    results = [[] for _ in values_lists]

    for rng in range(0, num_iterations, iteration_batch_size):
        max_rng = min(iteration_batch_size, num_iterations - rng)

        values_sims = _generate_distributions(values_lists, max_rng)

        for i, values_sim, stat_func in zip(range(len(values_sims)), values_sims, stat_func_lists):
            results[i].extend(stat_func(values_sim))

    return _np.array(results)

def _bootstrap_distribution(values_lists, stat_func_lists,
                            num_iterations, iteration_batch_size, num_threads):

    '''Returns the simulated bootstrap distribution. The idea is to sample the same
        indexes in a bootstrap re-sample across all arrays passed into values_lists.

        This is especially useful when you want to co-sample records in a ratio metric.
            numerator[k].sum() / denominator[k].sum()
        and not
            numerator[ j ].sum() / denominator[k].sum()
    Args:
        values_lists: list of numpy arrays (or scipy.sparse.csr_matrix)
            each represents a set of values to bootstrap. All arrays in values_lists
            must be of the same length.
        stat_func_lists: statistic to bootstrap for each element in values_lists.
        num_iterations: number of bootstrap iterations / resamples / simulations
            to perform.
        iteration_batch_size: The bootstrap sample can generate very large
            matrices. This argument limits the memory footprint by
            batching bootstrap rounds. If unspecified the underlying code
            will produce a matrix of len(values) x num_iterations. If specified
            the code will produce sets of len(values) x iteration_batch_size
            (one at a time) until num_iterations have been simulated.
            Defaults to no batching.
        num_threads: The number of therads to use. This speeds up calculation of
            the bootstrap. Defaults to 1. If -1 is specified then
            multiprocessing.cpu_count() is used instead.
    Returns:
        The set of bootstrap resamples where each stat_function is applied on
        the bootsrapped values.
    '''

    _validate_arrays(values_lists)

    if iteration_batch_size is None:
        iteration_batch_size = num_iterations

    num_iterations = int(num_iterations)
    iteration_batch_size = int(iteration_batch_size)

    num_threads = int(num_threads)

    if num_threads == -1:
        num_threads = _multiprocessing.cpu_count()

    if num_threads <= 1:
        results = _bootstrap_sim(values_lists, stat_func_lists,
                                 num_iterations, iteration_batch_size, None)
    else:
        pool = _multiprocessing.Pool(num_threads)

        iter_per_job = _np.ceil(num_iterations * 1.0 / num_threads)

        results = []
        for seed in _np.random.randint(0, 2**32 - 1, num_threads):
            r = pool.apply_async(_bootstrap_sim, (values_lists, stat_func_lists,
                                 iter_per_job,
                                 iteration_batch_size, seed))
            results.append(r)

        results = _np.hstack([res.get() for res in results])

        pool.close()

    return results

def bootstrap(values, stat_func, denominator_values=None, alpha=0.05,
              num_iterations=10000, iteration_batch_size=10, is_pivotal=True,
              num_threads=1, return_distribution=False):
    '''Returns bootstrap estimate.
    Args:
        values: numpy array (or scipy.sparse.csr_matrix) of values to bootstrap
        stat_func: statistic to bootstrap. We provide several default functions:
                * stat_functions.mean
                * stat_functions.sum
                * stat_functions.std
        denominator_values: optional array that does division after the
            statistic is aggregated. This lets you compute group level division
            statistics. One corresponding entry per record in @values.
            Example:
                SUM(value) / SUM(denom) instead of MEAN(value / denom)

                Ex. Cost Per Click
                cost per click across a group
                    SUM(revenue) / SUM(clicks)
                mean cost per click for each
                    MEAN(revenue / clicks)
        alpha: alpha value representing the confidence interval.
            Defaults to 0.05, i.e., 95th-CI.
        num_iterations: number of bootstrap iterations to run. The higher this
            number the more sure you can be about the stability your bootstrap.
            By this - we mean the returned interval should be consistent across
            runs for the same input. This also consumes more memory and makes
            analysis slower. Defaults to 10000.
        iteration_batch_size: The bootstrap sample can generate very large
            matrices. This argument limits the memory footprint by
            batching bootstrap rounds. If unspecified the underlying code
            will produce a matrix of len(values) x num_iterations. If specified
            the code will produce sets of len(values) x iteration_batch_size
            (one at a time) until num_iterations have been simulated.
            Defaults to 10. Passing None will calculate the full simulation in one step.
        is_pivotal: if true, use the pivotal method for bootstrapping confidence
            intervals. If false, use the percentile method.
        num_threads: The number of therads to use. This speeds up calculation of
            the bootstrap. Defaults to 1. If -1 is specified then
            multiprocessing.cpu_count() is used instead.
    Returns:
        BootstrapResults representing CI and estimated value.
    '''
    if denominator_values is None:
        values_lists = [values]
        stat_func_lists = [stat_func]

        def do_division(x):
            return x

        stat_val = stat_func(values)[0]
    else:
        values_lists = [values, denominator_values]
        stat_func_lists = [stat_func] * 2

        def do_division(num, denom):
            return num / denom

        stat_val = stat_func(values)[0] / stat_func(denominator_values)[0]

    distribution_results = _bootstrap_distribution(values_lists,
                                                   stat_func_lists,
                                                   num_iterations,
                                                   iteration_batch_size,
                                                   num_threads)

    bootstrap_dist = do_division(*distribution_results)

    if return_distribution:
        return bootstrap_dist
    else:
        return _get_confidence_interval(bootstrap_dist, stat_val, alpha,
                                        is_pivotal)

def bootstrap_ab(test, ctrl, stat_func, compare_func, test_denominator=None,
                 ctrl_denominator=None, alpha=0.05, num_iterations=10000,
                 iteration_batch_size=10, scale_test_by=1.0,
                 is_pivotal=True, num_threads=1, return_distribution=False):
    '''Returns bootstrap confidence intervals for an A/B test.
    Args:
        test: numpy array (or scipy.sparse.csr_matrix) of test results
        ctrl: numpy array (or scipy.sparse.csr_matrix) of ctrl results
        stat_func: statistic to bootstrap. We provide several default functions:
                * stat_functions.mean
                * stat_functions.sum
                * stat_functions.std
        compare_func: Function to compare test and control against.
                * compare_functions.difference
                * compare_functions.percent_change
                * compare_functions.ratio
                * compare_functions.percent_difference
        test_denominator: optional array that does division after the statistic
            is aggregated. This lets you compute group level division
            statistics. One corresponding entry per record in test.
            Example:
                SUM(value) / SUM(denom) instead of MEAN(value / denom)
                Ex. Cost Per Click
                cost per click across a group  (clicks is denominator)
                    SUM(revenue) / SUM(clicks)
                mean cost per click for each record
                    MEAN(revenue / clicks)
        ctrl_denominator: see test_denominator.
        alpha: alpha value representing the confidence interval.
            Defaults to 0.05, i.e., 95th-CI.
        num_iterations: number of bootstrap iterations to run. The higher this
            number the more sure you can be about the stability your bootstrap.
            By this - we mean the returned interval should be consistent across
            runs for the same input. This also consumes more memory and makes
            analysis slower.
        iteration_batch_size: The bootstrap sample can generate very large
            arrays. This function iteration_batch_size limits the memory
            footprint by batching bootstrap rounds. Defaults to 10. Passing None
            will attempt to calculate the full simulation in one step.
        scale_test_by: The ratio between test and control population
            sizes. Use this if your test and control split is different from a
            50/50 split. Defaults to 1.0.
        is_pivotal: if true, use the pivotal method for bootstrapping confidence
            intervals. If false, use the percentile method.
        num_threads: The number of therads to use. This speeds up calculation of
            the bootstrap. Defaults to 1. If -1 is specified then
            multiprocessing.cpu_count() is used instead.
    Returns:
        BootstrapResults representing CI and estimated value.
    '''

    both_denominators = test_denominator is not None and \
            ctrl_denominator is not None
    both_numerators = test is not None and ctrl is not None

    if both_numerators and not both_denominators:
        test_lists = [test]
        ctrl_lists = [ctrl]
        stat_func_lists = [stat_func]

        def do_division(x):
            return x

        test_val = stat_func(test)[0]
        ctrl_val = stat_func(ctrl)[0]

    elif both_numerators and both_denominators:
        test_lists = [test, test_denominator]
        ctrl_lists = [ctrl, ctrl_denominator]
        stat_func_lists = [stat_func] * 2

        def do_division(num, denom):
            return num / denom

        test_val = stat_func(test)[0] / stat_func(test_denominator)[0]
        ctrl_val = stat_func(ctrl)[0] / stat_func(ctrl_denominator)[0]

    elif not both_numerators:
        raise ValueError('Both test and ctrl numerators must be specified.')
    else:
        raise ValueError('Both test and ctrl denominators must be specified.')

    test_results = _bootstrap_distribution(test_lists, stat_func_lists,
                                           num_iterations, iteration_batch_size,
                                           num_threads)

    ctrl_results = _bootstrap_distribution(ctrl_lists, stat_func_lists,
                                           num_iterations, iteration_batch_size,
                                           num_threads)

    test_dist = do_division(*test_results)
    ctrl_dist = do_division(*ctrl_results)

    test_ctrl_dist = compare_func(test_dist * scale_test_by, ctrl_dist)

    if return_distribution:
        return test_ctrl_dist
    else:

        test_ctrl_val = compare_func(test_val * scale_test_by, ctrl_val)
        return _get_confidence_interval(test_ctrl_dist, test_ctrl_val, alpha,
                                        is_pivotal)


#! /user/bin/env python
# encoding:utf-8
import requests
import pandas as pd
import json

markets = ['szmb','szsme','szcn','shmb']
market_names = ['深市主板','中小板','创业板','沪市']
market_name_dict = dict(zip(markets, market_names))

def fetch_prbookinfo_from_cninfo(market):
    #从巨潮资讯获取报告披露时间信息
    #market是市场代码
    print('开始获取'+market_name_dict[market]+'数据.......')
    base_url = 'http://www.cninfo.com.cn/cninfo-new/information/prbookinfo-1'
    data = {'sectionTime': '2016-12-31', 'market': market, 'orderClos': '', 'isDesc': '', 'stockCode': '',
            'firstTime': u'请输入预约披露时间', 'lastTime': '请输入实际披露时间'}
    r = requests.post(base_url,data=data)
    if r.status_code != 200:
        raise Exception('Error in get data,status_code is:'+str(r.status_code))
    datas = json.loads(r.text)
    df = pd.DataFrame(datas)
    print('共获取'+market_name_dict[market]+'股票 '+ str(len(df)) + '只')
    return df

def fetch_no_annualReport_stocks():
    # 筛选没有出年报的股票
    df = None
    for market in markets:
        t_df = fetch_prbookinfo_from_cninfo(market)
        #实际披露时间形如2017-03-17，至少有10个字符
        t_df = t_df[t_df.f006d_0102.str.len() < 10]
        print(market_name_dict[market]+'尚未披露年报的股票有' + str(len(t_df)) + '只')
        if df is None:
            df = t_df
        else:
            df = df.append(t_df,ignore_index=True)
    print('还没有披露年报的股票一共有 '+ str(len(df))+ '只')
    return df

def main():
    hdf = fetch_no_annualReport_stocks()
    print(hdf)
    # filename = u'F:/firecapital/数据/output20170501/普通总列表.xlsx'
    # df = pd.read_excel(filename,converters = {'code':str})
    # df = df[['period','code','weight']]
    # if len(df.code[0]) > 6:
    #     df.code = df.code.str.slice(2,8)

if __name__ == '__main__':
    main()


import asyncio
import bz2
import contextlib
import logging

import player.codec
import player.terminal
import player.tools

class Session:
	'''
	Session represent one client
	also separate sessions can replay different files
	'''

	def __init__(self, session_id: str, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, filename: str):
		self._logger = logging.getLogger('{}[{}]'.format(self.__class__.__name__, session_id))
		self.session_id = session_id
		self.filename = filename
		self._writer = writer
		self._reader = reader

		self._run_future = None # type: asyncio.Future

	async def run(self) -> None:
		'''
		Wrapper which allow us to terminate session object
		'''
		self._run_future = asyncio.ensure_future(self._run())
		await self._run_future

	async def _run(self) -> None:
		'''
		Run loop which will stream frame by frame to writer
		'''
		self._logger.debug('Reading file = %s', self.filename)
		with bz2.BZ2File(self.filename, 'rb') as file:
			await self._clear_screen()

			metadata = player.codec.get_file_metadata(file)
			sleep_time = 1.0 / metadata.frame_rate

			stop_watch = player.tools.StopWatch()
			for frame in player.codec.get_frames(file, metadata):
				# Move cursor to top left corner
				self._writer.write(player.terminal.RESET_CURSOR)
				self._writer.write(frame)
				await self._writer.drain()

				# We will have to correct sleep time and
				# decrease how long it take to load one frame from it
				t = max(0, stop_watch.lap() - sleep_time)
				if sleep_time < t:
					self._logger.error('Sleep time is less then reading overhead by = %.2f', t)
				await asyncio.sleep(max(0, sleep_time - t))

			await self._clear_screen()

	async def terminate(self) -> None:
		'''
		Close writer and clenaup after session
		'''
		self._logger.debug('Terminating of session requested.')
		if self._run_future is not None and not self._run_future.done():
			self._run_future.cancel()

			with contextlib.suppress(asyncio.CancelledError):
				await self._run_future

		self._writer.close()

	async def _clear_screen(self) -> None:
		'''
		Write clearing of screen
		'''
		# First white all empty lines
		self._writer.write(player.terminal.CLEAR)
		# And then reset cursor to upper left
		self._writer.write(player.terminal.RESET_CURSOR)
		await self._writer.drain()


from typing import List, Optional, Sequence

from llama_index.llms.types import ChatMessage, MessageRole

BOS, EOS = "<s>", "</s>"
B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = """\
You are a helpful, respectful and honest assistant. \
Always answer as helpfully as possible and follow ALL given instructions. \
Do not speculate or make up information. \
Do not reference any given instructions or context. \
"""

def messages_to_prompt(
    messages: Sequence[ChatMessage], system_prompt: Optional[str] = None
) -> str:
    string_messages: List[str] = []
    if messages[0].role == MessageRole.SYSTEM:
        # pull out the system message (if it exists in messages)
        system_message_str = messages[0].content or ""
        messages = messages[1:]
    else:
        system_message_str = system_prompt or DEFAULT_SYSTEM_PROMPT

    system_message_str = f"{B_SYS} {system_message_str.strip()} {E_SYS}"

    for i in range(0, len(messages), 2):
        # first message should always be a user
        user_message = messages[i]
        assert user_message.role == MessageRole.USER

        if i == 0:
            # make sure system prompt is included at the start
            str_message = f"{BOS} {B_INST} {system_message_str} "
        else:
            # end previous user-assistant interaction
            string_messages[-1] += f" {EOS}"
            # no need to include system prompt
            str_message = f"{BOS} {B_INST} "

        # include user message content
        str_message += f"{user_message.content} {E_INST}"

        if len(messages) > (i + 1):
            # if assistant message exists, add to str_message
            assistant_message = messages[i + 1]
            assert assistant_message.role == MessageRole.ASSISTANT
            str_message += f" {assistant_message.content}"

        string_messages.append(str_message)

    return "".join(string_messages)

def completion_to_prompt(completion: str, system_prompt: Optional[str] = None) -> str:
    system_prompt_str = system_prompt or DEFAULT_SYSTEM_PROMPT

    return (
        f"{BOS} {B_INST} {B_SYS} {system_prompt_str.strip()} {E_SYS} "
        f"{completion.strip()} {E_INST}"
    )


# -*- coding: utf-8 -*-
from __future__ import unicode_literals
from django.http import HttpResponse
from django.shortcuts import render
from django.views.generic import TemplateView
from panampath.models import PathSegment
from pprint import pprint
from forms import EmailForm
from django.conf import settings
from django.shortcuts import render_to_response, render
from django.core.mail.message import EmailMessage

# Create your views here.
class IndexView(TemplateView):
    template_name = 'main.html'

    def get_context_data(self, **kwargs):
        context = super(IndexView, self).get_context_data(**kwargs)
        segments = PathSegment.objects.all()

        coords = []
        for segment in segments:
            if segment.geom:
                coords.append(segment.geom['coordinates'])
            # for member in segment.geom['coordinates']:
            #     for item in member:
            #         coords.append(item[::-1])

        context["segments"] = segments
        context["coordinates"] = coords

        return context

def send_email(request):

    if request.method != 'POST':
        form = EmailForm()
        return render(request, 'submit_event.html', {'email_form': form})

    form = EmailForm(request.POST, request.FILES)

    if form.is_valid():
        subject = form.cleaned_data['subject']
        message = form.cleaned_data['message']
        email = form.cleaned_data['email']
        attach = request.FILES['attach']
        try:
            mail = EmailMessage(subject, message, settings.EMAIL_HOST_USER, [email])
            mail.attach(attach.name, attach.read(), attach.content_type)
            response = mail.send()
            return render(request, 'submit_event.html', {'message': 'Sent email to %s' % (email)})
        except Exception as e:
            return render(request, 'submit_event.html', {'message': e.message})
    else:
        try:
            return render(request, 'submit_event.html', {'message': form.message})

        except AttributeError:
            return render(request, 'submit_event.html', {'message': "Please fill out all fields on the form.", "email_form":form})

    #return render(request, 'submit_event.html', {'message': 'Unable to send email. Please try again later'})


import json
import logging
import configparser
import subprocess
import sys

ALIAS = sys.argv[1]
RESOURCE_ROLE = sys.argv[2]

cfg = configparser.ConfigParser()
cfg.read('deployment_config.ini')

# parse lambda configuration parameters
CONFIG_PARAMS = {}
for config_tuple in cfg.items('configuration'):
    name = config_tuple[0]
    val = config_tuple[1]
    CONFIG_PARAMS[name] = val

# https://stackoverflow.com/questions/5466451/how-can-i-print-literal-curly-brace-characters-in-python-string-and-also-use-fo#5466478
VARS_STR_TEMPLATE = "Variables={{{}}}"
env_vars = ""
for config_tuple in cfg.items('environment'):
    env_vars += "{k}={v},".format(k=config_tuple[0], v=config_tuple[1])

if env_vars.endswith(","):
    env_vars = env_vars[:-1]
VARS_STR = VARS_STR_TEMPLATE.format(env_vars)
print("Parsed environment variables:\n{}".format(VARS_STR))

def is_func_new(funcname):
    """
    determine if function being deployed is brand new or just needs updates
    """
    bashCommand = "aws lambda get-function \
    --function-name {fname}".format(
        fname=funcname
    )

    try:
        subprocess.check_output(bashCommand.split())
    except subprocess.CalledProcessError as e:
        # if the error is raised it means the functions does not exist
        logging.error(e)
        print("returning True")
        return True
    print("returning False")
    return False

def deploy_lambda(new):
    """
    ISNEW=$1
    RESOURCE_ROLE=$2
    ALIAS=$3

    FNAME=$4
    HANDLER=$5
    TIMEOUT=$6
    MEMSIZE=$7
    DESC=$8
    ENV=$9
    RUNTIME=${10}
    REGION=${11}
    """
    bashCommand = "bash scripts/aws-lambda-deploy.sh \
    {isnew} \
    {AWS_LAMBDA_ROLE} \
    {alias} \
    {funcname} {handler} {timeout} {memsize} '{desc}' '{env}' '{runtime}' '{region}'".format(
        AWS_LAMBDA_ROLE=RESOURCE_ROLE,
        isnew=new,
        alias=ALIAS,
        funcname=CONFIG_PARAMS['function-name'],
        desc=CONFIG_PARAMS['description'],
        runtime=CONFIG_PARAMS['runtime'],
        handler=CONFIG_PARAMS['handler'],
        region=CONFIG_PARAMS['region'],
        timeout=CONFIG_PARAMS['timeout'],
        memsize=CONFIG_PARAMS['memory-size'],
        env=VARS_STR
    )

    process = subprocess.Popen(bashCommand, stdout=subprocess.PIPE, shell=True)
    output, error = process.communicate()
    # print(output)
    if error:
        logging.error(error)

deploy_lambda(new=is_func_new(CONFIG_PARAMS['function-name']))


import os
import logging
from logging.handlers import SMTPHandler, RotatingFileHandler
import pprint

from elasticsearch import Elasticsearch
from flask import Flask, session, request, current_app, redirect
from flask_sqlalchemy import SQLAlchemy
from flask_migrate import Migrate
from flask_login import LoginManager, current_user
from flask_mail import Mail
from flask_moment import Moment
from flask_babel import Babel
from flask_admin import Admin, BaseView, expose
from flask_admin.contrib.sqla import ModelView

from config import Config
from app.search import get_mappings, insert_mapping, delete_mapping

# Turn off autoflush to let review editing to be saved in session.dirty
db = SQLAlchemy(session_options={"autoflush": False})
migrate = Migrate()
login = LoginManager()
login.login_view = 'auth.login'
mail = Mail()
moment = Moment()
babel = Babel()
admin = Admin()

def create_app(config_class=Config):
    app = Flask(__name__)
    app.config.from_object(config_class)

    db.init_app(app)
    migrate.init_app(app, db)
    login.init_app(app)
    mail.init_app(app)
    moment.init_app(app)
    babel.init_app(app)
    admin.init_app(app)
    app.elasticsearch = Elasticsearch([app.config['ELASTICSEARCH_URL']]) if app.config['ELASTICSEARCH_URL'] else None

    from app.errors.handlers import bp as errors_bp
    app.register_blueprint(errors_bp)

    from app.auth import bp as auth_bp
    app.register_blueprint(auth_bp, url_prefix='/auth')

    from app.main import bp as main_bp
    app.register_blueprint(main_bp)

    if not app.debug and not app.testing:
        if app.config['MAIL_SERVER']:
            auth = None
            if app.config['MAIL_USERNAME'] or app.config['MAIL_PASSWORD']:
                auth = (app.config['MAIL_USERNAME'], app.config['MAIL_PASSWORD'])
            secure = None
            if app.config['MAIL_USE_TLS']:
                secure = ()
            mail_handler = SMTPHandler(
                mailhost=(app.config['MAIL_SERVER'], app.config['MAIL_PORT']),
                fromaddr='no-reply@' + app.config['MAIL_SERVER'],
                toaddrs=app.config['ADMINS'], subject='Whisky Blog failure',
                credentials=auth, secure=secure)
            mail_handler.setLevel(logging.ERROR)
            app.logger.addHandler(mail_handler)
        if app.config['LOG_TO_STDOUT']:
            stream_handler = logging.StreamHandler()
            stream_handler.setLevel(logging.INFO)
            app.logger.addHandler(stream_handler)
        else:
            if not os.path.exists('logs'):
                os.mkdir('logs')
            file_handler = RotatingFileHandler('logs/WhiskyBlog.logs', maxBytes=10240, backupCount=10)
            file_handler.setFormatter(logging.Formatter(
                '%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]'))
            file_handler.setLevel(logging.INFO)
            app.logger.addHandler(file_handler)

        app.logger.setLevel(logging.INFO)
        app.logger.info('Whisky Blog startup')

    return app

# Looks for the language the current flask app is set to.
@babel.localeselector
def get_locale():
    try:
        language = session['language']
    except KeyError:
        language = None
    if language is not None:
        return language
    return request.accept_languages.best_match(current_app.config['LANGUAGES'])

from app import models

"""Create custom admin views for `User`, `Review` and `Tag` models"""

class BaseModelView(ModelView):
    def is_accessible(self):
        return current_user.is_authenticated and current_user.id == 1

class UserView(BaseModelView):
    column_exclude_list = ['password_hash']
    can_create = False
    can_edit = True

class ReviewView(BaseModelView):
    can_create = False
    can_edit = False

class TagView(BaseModelView):
    @expose('/bulk/')
    def bulk_view(self):
        from app.main.info import all_tags
        for t in all_tags:
            tag = self.model.query.filter_by(name=t[0]).first()
            if not tag:
                self.session.add(self.model(name=t[0]))
        self.session.commit()
        return redirect('/admin/tag')

class SearchView(BaseView):
    @expose('/')
    def index(self):
        maps = get_mappings()
        return self.render('admin/search.html', mapping=pprint.pformat(maps))

    @expose('/insert/')
    def insert(self):
        # insert elasticsearch mapping
        insert_mapping('review')
        return redirect('/admin/search')

    @expose('/delete/')
    def delete(self):
        # delete elasticsearch mapping
        delete_mapping('review')
        return redirect('/admin/search')

    def is_accessible(self):
        return current_user.is_authenticated and current_user.id == 1

admin.add_view(UserView(models.User, db.session))
admin.add_view(ReviewView(models.Review, db.session))
admin.add_view(TagView(models.Tag, db.session))
admin.add_view(SearchView(name='Search', endpoint='search'))


from __future__ import annotations

import logging

from enum import Enum, unique
from typing import TYPE_CHECKING, List

import numpy as np

from scipy.linalg import block_diag
from scipy.special import digamma, gammaln

from stmmap.utils.distances import DistanceMetric, distance_gauss, distance_invgamma

if TYPE_CHECKING:
    from stmmap.surfel import Surfel

logger = logging.getLogger(__name__)

class ProjectionMethod:
    def __init__(self, surf: Surfel):
        self.surfel = surf
        self.dim = surf.dim

        self.iterations = 0

        self._load()

    def _load(self):
        pass

    def iteration(self) -> float:
        """Iterate once through all the measurements."""
        pot_old_h_C = np.copy(self.surfel.pot_h_C)
        pot_old_h_m = np.copy(self.surfel.pot_h_m)

        bel_old_v_a = np.copy(self.surfel.bel_v_a)
        bel_old_v_b = np.copy(self.surfel.bel_v_b)

        batch_size = self.surfel.n_window

        if batch_size == 0:
            return 0.0

        order = np.random.choice(range(batch_size), batch_size, replace=False)
        for i in order:
            self.surfel.pot_h_h -= self.surfel.like_outmsg_h_h[i]
            self.surfel.pot_h_K -= self.surfel.like_outmsg_h_K[i]

            self.update_likelihood_cluster_potential(i)

            self.surfel.pot_h_h += self.surfel.like_outmsg_h_h[i]
            self.surfel.pot_h_K += self.surfel.like_outmsg_h_K[i]

        # self.surfel.pot_h_h = self.surfel.bel_h_h - self.surfel.inmsg_h_h
        # self.surfel.pot_h_K = self.surfel.bel_h_K - self.surfel.inmsg_h_K
        # # self.surfel.pot_h_h = self.surfel.prior_h_h + np.sum(
        # #     self.surfel.like_outmsg_h_h[: self.surfel.n_window], axis=0
        # # )
        # # self.surfel.pot_h_K = self.surfel.prior_h_K + np.sum(
        # #     self.surfel.like_outmsg_h_K[: self.surfel.n_window], axis=0
        # # )
        self.surfel.pot_h_C = np.linalg.inv(self.surfel.pot_h_K)
        self.surfel.pot_h_m = self.surfel.pot_h_C.dot(self.surfel.pot_h_h)

        # self.surfel.bel_h_h = self.surfel.pot_h_h + self.surfel.inmsg_h_h
        # self.surfel.bel_h_K = self.surfel.pot_h_K + self.surfel.inmsg_h_K
        # self.surfel.bel_h_C = np.linalg.inv(self.surfel.bel_h_K)
        # self.surfel.bel_h_m = self.surfel.bel_h_C.dot(self.surfel.bel_h_h)
        # distance = distance_gauss(
        #     bel_old_h_m, bel_old_h_C, self.surfel.bel_h_m, self.surfel.bel_h_C
        # )

        distance = distance_gauss(
            pot_old_h_m, pot_old_h_C, self.surfel.pot_h_m, self.surfel.pot_h_C
        )
        distance += distance_invgamma(
            bel_old_v_a, bel_old_v_b, self.surfel.bel_v_a, self.surfel.bel_v_b
        )

        return distance

    def update_likelihood_cluster_potential(self, i: int) -> None:
        """Update the likelihood cluster potential for the i-th measurement.

        Note this should also update the surfel belief.

        """
        raise NotImplementedError()

class VMPStructured(ProjectionMethod):
    def _load(self):
        # Predefine some matrices for the EIF
        taylor_dims = 2 * self.dim - 1
        self.G = np.zeros((2 * self.dim, taylor_dims), dtype=float)
        self.G[:taylor_dims, :] = np.eye(taylor_dims, dtype=float)

        self.C_stack = 100 * np.eye(taylor_dims, dtype=float)  # h, a, b

        # Selection matix
        self.S = np.zeros((self.dim, 2 * self.dim), dtype=float)
        self.S[:, self.dim :] = np.eye(self.dim, dtype=float)

    def update_likelihood_cluster_potential(self, i: int) -> None:
        logging_pre = f"z[{i}] : "

        m_z = self.surfel.z_m[i, :, 0]

        # Update heights
        # [ Calculate cavity distribution
        self.surfel.bel_h_K -= self.surfel.like_outmsg_h_K[i]
        self.surfel.bel_h_h -= self.surfel.like_outmsg_h_h[i]

        h_cav_K = np.copy(self.surfel.bel_h_K)
        h_cav_h = np.copy(self.surfel.bel_h_h)
        h_cav_C = np.linalg.inv(h_cav_K)
        h_cav_m = h_cav_C.dot(h_cav_h)
        # ]

        # EIF Prediction Update
        # [ Taylor
        m_pred = np.zeros((2 * self.dim, 1), dtype=float)
        m_pred[: self.dim] = h_cav_m
        a_prior = m_z[0]
        if self.dim == 3:
            b_prior = m_z[1]
            A = np.array((1 - a_prior - b_prior, a_prior, b_prior)).reshape(1, self.dim)
            m_pred[self.dim + 1] = b_prior
        else:  # 2-D
            A = np.array((1 - a_prior, a_prior)).reshape(1, self.dim)

        g_pred = A.dot(h_cav_m)
        m_pred[self.dim] = a_prior
        m_pred[-1] = g_pred

        self.G[-1, : self.dim] = A.flat
        self.G[-1, self.dim :] = h_cav_m[1:, 0] - h_cav_m[0, 0]

        self.C_stack[: self.dim, : self.dim] = h_cav_C
        C_pred = self.G.dot(self.C_stack).dot(self.G.T)
        C_pred[-1, -1] += self.surfel.bel_v_b / self.surfel.bel_v_a

        # Numerical Symmetry
        C_pred += C_pred.T
        C_pred /= 2

        if (np.diag(C_pred) < 0).any():
            msg = (
                logging_pre + "Weird prediction covariance matrix\n"
                f"Rank: {np.linalg.matrix_rank(C_pred)}\n"
                f"Diagonal: {np.diag(C_pred)}"
            )
            logger.warn(msg)

        # ]

        # EIF Measurement Update
        K_pred = np.linalg.inv(C_pred)
        h_pred = K_pred.dot(m_pred)
        z = m_z.reshape(self.dim, 1)
        K_tmp = K_pred + self.S.T.dot(self.surfel.z_K[i]).dot(self.S)
        h_tmp = h_pred + self.S.T.dot(self.surfel.z_K[i]).dot(z)

        # Numerical Symmetry
        K_tmp += K_tmp.T
        K_tmp /= 2

        C_tmp = np.linalg.inv(K_tmp)
        m_tmp = C_tmp.dot(h_tmp)

        if (np.diag(C_tmp) < 0).any():
            msg = (
                logging_pre + "Weird height belief covariance matrix\n"
                f"Rank: {np.linalg.matrix_rank(C_tmp)}\n"
                f"Diagonal: {np.diag(C_tmp)}"
            )
            logger.warn(msg)

        # Calculate measurement factor
        m_C_tmp = np.linalg.inv(K_tmp[self.dim :, self.dim :])
        self.surfel.like_outmsg_h_K[i] = (
            K_tmp[: self.dim, : self.dim]
            - self.surfel.bel_h_K
            - K_tmp[: self.dim, self.dim :].dot(m_C_tmp).dot(K_tmp[self.dim :, : self.dim])
        )
        self.surfel.like_outmsg_h_h[i] = (
            h_tmp[: self.dim, :]
            - self.surfel.bel_h_h
            - K_tmp[: self.dim, self.dim :].dot(m_C_tmp).dot(h_tmp[self.dim :, :])
        )

        # Update posterior
        self.surfel.bel_h_K += self.surfel.like_outmsg_h_K[i]
        self.surfel.bel_h_h += self.surfel.like_outmsg_h_h[i]
        self.surfel.bel_h_C = np.linalg.inv(self.surfel.bel_h_K)
        self.surfel.bel_h_m = self.surfel.bel_h_C.dot(self.surfel.bel_h_h)

        # Update variation
        # [ Calculate cavity distribution
        self.surfel.bel_v_a -= self.surfel.like_outmsg_v_a[i]
        self.surfel.bel_v_b -= self.surfel.like_outmsg_v_b[i]

        variation = self.surfel.bel_v_b / self.surfel.bel_v_a
        if variation < 0:
            msg = (
                logging_pre
                + f"The deviation cavity distribution expected value is negative: {variation}"
            )
            logger.warn(msg)
        # ]

        # expectation = self._exact_expectation(m_tmp, C_tmp)
        expectation = self._taylor_expectation(m_tmp, C_tmp)  # This is faster

        self.surfel.like_outmsg_v_b[i] = expectation / 2.0
        self.surfel.like_outmsg_v_a[i] = 0.5

        self.surfel.bel_v_a += self.surfel.like_outmsg_v_a[i]
        self.surfel.bel_v_b += self.surfel.like_outmsg_v_b[i]
        variation = self.surfel.bel_v_b / self.surfel.bel_v_a
        if variation < 0:
            msg = logging_pre + f"The deviation belief expected value is negative: {variation}"
            logger.warn(msg)

    def _exact_expectation(self, m, C):
        """Exact expectation calculation.

        This was derived using sympy.

        E_g(h,m)[(gamma - f(alpha, h))^2]

        """

        def second_order(m, C):
            return (m[0,] * m[1,] + C[0, 1]).item()

        def forth_order(m, C):
            E = 0.0
            E += second_order(m[(0, 1),], C[(0, 1),][:, (0, 1)]) * second_order(
                m[(2, 3),], C[(2, 3),][:, (2, 3)]
            )
            E += second_order(m[(0, 2),], C[(0, 2),][:, (0, 2)]) * second_order(
                m[(1, 3),], C[(1, 3),][:, (1, 3)]
            )
            E += second_order(m[(0, 3),], C[(0, 3),][:, (0, 3)]) * second_order(
                m[(1, 2),], C[(1, 2),][:, (1, 2)]
            )
            E += -2 * m[0,] * m[1,] * m[2,] * m[3,]

            return E.item()

        def third_order(m, C):
            E = 0.0
            E += m[0,] * m[1,] * m[2,]
            E += m[0,] * C[1, 2]
            E += m[1,] * C[0, 2]
            E += m[2,] * C[0, 1]

            return E.item()

        if self.dim == 3:
            # h0 ha hb a  b  g
            # 0  1  2  3  4  5
            E = 0.0
            # XXX 4th order
            # a**2*h0**2     | +1 | 0 0 3 3
            indices = (0, 0, 3, 3)
            E += forth_order(m[indices,], C[indices,][:, indices])
            # - 2*a**2*h0*ha | -2 | 0 1 3 3
            indices = (0, 1, 3, 3)
            E += -2 * forth_order(m[indices,], C[indices,][:, indices])
            # + a**2*ha**2   | +1 | 1 1 3 3
            indices = (1, 1, 3, 3)
            E += forth_order(m[indices,], C[indices,][:, indices])
            # + 2*a*b*h0**2  | +2 | 0 0 3 4
            indices = (0, 0, 3, 4)
            E += 2 * forth_order(m[indices,], C[indices,][:, indices])
            # - 2*a*b*h0*ha  | -2 | 0 1 3 4
            indices = (0, 1, 3, 4)
            E += -2 * forth_order(m[indices,], C[indices,][:, indices])
            # - 2*a*b*h0*hb  | -2 | 0 2 3 4
            indices = (0, 2, 3, 4)
            E += -2 * forth_order(m[indices,], C[indices,][:, indices])
            # + 2*a*b*ha*hb  | +2 | 1 2 3 4
            indices = (1, 2, 3, 4)
            E += 2 * forth_order(m[indices,], C[indices,][:, indices])
            # + b**2*h0**2   | +1 | 0 0 4 4
            indices = (0, 0, 4, 4)
            E += forth_order(m[indices,], C[indices,][:, indices])
            # - 2*b**2*h0*hb | -2 | 0 2 4 4
            indices = (0, 2, 4, 4)
            E += -2 * forth_order(m[indices,], C[indices,][:, indices])
            # + b**2*hb**2   | +1 | 2 2 4 4
            indices = (2, 2, 4, 4)
            E += forth_order(m[indices,], C[indices,][:, indices])

            # XXX 3rd order
            # + 2*a*y*h0     | +2 | 0 3 5
            indices = (0, 3, 5)
            E += 2 * third_order(m[indices,], C[indices,][:, indices])
            # - 2*a*y*ha     | -2 | 1 3 5
            indices = (1, 3, 5)
            E += -2 * third_order(m[indices,], C[indices,][:, indices])
            # - 2*a*h0**2    | -2 | 0 0 3
            indices = (0, 0, 3)
            E += -2 * third_order(m[indices,], C[indices,][:, indices])
            # + 2*a*h0*ha    | +2 | 0 1 3
            indices = (0, 1, 3)
            E += 2 * third_order(m[indices,], C[indices,][:, indices])
            # + 2*b*y*h0     | +2 | 0 4 5
            indices = (0, 4, 5)
            E += 2 * third_order(m[indices,], C[indices,][:, indices])
            # - 2*b*y*hb     | -2 | 2 4 5
            indices = (2, 4, 5)
            E += -2 * third_order(m[indices,], C[indices,][:, indices])
            # - 2*b*h0**2    | -2 | 0 0 4
            indices = (0, 0, 4)
            E += -2 * third_order(m[indices,], C[indices,][:, indices])
            # + 2*b*h0*hb    | +2 | 0 2 4
            indices = (0, 2, 4)
            E += 2 * third_order(m[indices,], C[indices,][:, indices])

            # XXX 2nd order
            # + y**2         | +1 | 5 5
            indices = (5, 5)
            E += second_order(m[indices,], C[indices,][:, indices])
            # - 2*y*h0       | -2 | 0 5
            indices = (0, 5)
            E += -2 * second_order(m[indices,], C[indices,][:, indices])
            # + h0**2        | +1 | 0 0
            indices = (0, 0)
            E += second_order(m[indices,], C[indices,][:, indices])
        else:  # 2-D
            # h0 ha a g
            # 0  1  2 3
            E = 0.0
            # XXX 4th order
            # a**2*h0**2     | +1 | 0 0 2 2
            indices = (0, 0, 2, 2)
            E += forth_order(m[indices,], C[indices,][:, indices])
            # - 2*a**2*h0*ha | -2 | 0 1 2 2
            indices = (0, 1, 2, 2)
            E += -2 * forth_order(m[indices,], C[indices,][:, indices])
            # + a**2*ha**2   | +1 | 1 1 2 2
            indices = (1, 1, 2, 2)
            E += forth_order(m[indices,], C[indices,][:, indices])

            # XXX 3rd order
            # + 2*a*g*h0     | +2 | 0 2 3
            indices = (0, 2, 3)
            E += 2 * third_order(m[indices,], C[indices,][:, indices])
            # - 2*a*g*ha     | -2 | 1 2 3
            indices = (1, 2, 3)
            E += -2 * third_order(m[indices,], C[indices,][:, indices])
            # - 2*a*h0**2    | -2 | 0 0 2
            indices = (0, 0, 2)
            E += -2 * third_order(m[indices,], C[indices,][:, indices])
            # + 2*a*h0*ha    | +2 | 0 1 2
            indices = (0, 1, 2)
            E += 2 * third_order(m[indices,], C[indices,][:, indices])

            # XXX 2nd order
            # + g**2         | +1 | 3 3
            indices = (3, 3)
            E += second_order(m[indices,], C[indices,][:, indices])
            # - 2*g*h0       | -2 | 0 3
            indices = (0, 3)
            E += -2 * second_order(m[indices,], C[indices,][:, indices])
            # + h0**2        | +1 | 0 0
            indices = (0, 0)
            E += second_order(m[indices,], C[indices,][:, indices])

        return E

    def _monte_expectation(self, m, C, N_samples=10000000):
        x = np.random.multivariate_normal(m[:, 0], C, N_samples).T

        h0 = x[0, :]
        ha = x[1, :]
        alphas = x[2, :]
        gammas = x[3, :]

        E = np.mean((gammas - ((ha - h0) * alphas + h0)) ** 2)

        return E

    def _taylor_expectation(self, m, C):
        if self.dim == 3:
            h_m = m[:3, :]
            m_m = m[3:, :]
            H = np.array(
                (
                    1 - m_m[0, 0] - m_m[1, 0],
                    m_m[0, 0],
                    m_m[1, 0],
                    h_m[1, 0] - h_m[0, 0],
                    h_m[2, 0] - h_m[0, 0],
                    1,
                )
            ).reshape(1, 6)
            g_pred = (
                (1 - m_m[0, 0] - m_m[1, 0]) * h_m[0, 0]
                + m_m[0, 0] * h_m[1, 0]
                + m_m[1, 0] * h_m[2, 0]
            )
            E_taylor = H.dot(C).dot(H.T) + (m_m[2] - g_pred) ** 2
        else:  # 2-D
            h_m = m[:2, :]
            m_m = m[2:, :]
            H = np.array((1 - m_m[0, 0], m_m[0, 0], h_m[1, 0] - h_m[0, 0], 1)).reshape(1, 4)
            g_pred = (1 - m_m[0, 0]) * h_m[0, 0] + m_m[0, 0] * h_m[1, 0]
            E_taylor = H.dot(C).dot(H.T) + (m_m[1] - g_pred) ** 2

        return E_taylor.item()

class VMPFactorised(ProjectionMethod):
    def update_likelihood_cluster_potential(self, i: int) -> None:
        # [ NOTE these are views not copies!
        z_m = self.surfel.z_m[i, :, :]
        z_h = self.surfel.z_h[i, :, :]
        z_K = self.surfel.z_K[i, :, :]
        # ]

        # Update q(m_i) ------------------------------------------------
        a_v = 100
        m_h_prior = np.zeros((self.dim, 1), dtype=float)
        m_h_prior[:-1] = z_m[:-1] / a_v
        m_K_prior = np.eye(self.dim, dtype=float) / a_v
        m_K_prior[-1, -1] = 0
        if self.dim == 3:
            B_m = np.array([[-1, -1, 0], [1, 0, 0], [0, 1, 0]], dtype=float)
            b_m = np.array([[1], [0], [0]], dtype=float)
            a_h = np.array(
                [
                    [self.surfel.bel_h_m[0, 0] - self.surfel.bel_h_m[1, 0]],
                    [self.surfel.bel_h_m[0, 0] - self.surfel.bel_h_m[2, 0]],
                    [1],
                ],
                dtype=float,
            )
        else:
            B_m = np.array([[-1, 0], [1, 0]], dtype=float)
            b_m = np.array([[1], [0]], dtype=float)
            a_h = np.array(
                [[self.surfel.bel_h_m[0, 0] - self.surfel.bel_h_m[1, 0]], [1]], dtype=float
            )
        m_K_like = B_m.T.dot(self.surfel.bel_h_C).dot(B_m)
        m_K_like += a_h.dot(a_h.T)
        m_h_like = -B_m.T.dot(self.surfel.bel_h_C).dot(b_m)
        tmp = np.zeros((self.dim, 1), dtype=float)
        tmp[-1, 0] = self.surfel.bel_h_m[0, 0]
        m_h_like += a_h.dot(a_h.T).dot(tmp)

        variation = self.surfel.bel_v_b / self.surfel.bel_v_a
        m_K_like /= variation
        m_h_like /= variation

        m_h = z_h + m_h_like + m_h_prior
        m_K = z_K + m_K_like + m_K_prior

        m_C = np.linalg.inv(m_K)
        m_m = m_C.dot(m_h)

        # Update q(h) --------------------------------------------------
        if self.dim == 3:
            B_h = np.array([[1, -1, 0], [1, 0, -1], [0, 0, 0]], dtype=float)
            b_h = np.array([[0], [0], [1]], dtype=float)
            a_m = np.array([[1 - m_m[0, 0] - m_m[1, 0]], [m_m[0, 0]], [m_m[1, 0]]], dtype=float)
        else:
            B_h = np.array([[1, -1], [0, 0]], dtype=float)
            b_h = np.array([[0], [1]], dtype=float)
            a_m = np.array([[1 - m_m[0, 0]], [m_m[0, 0]]], dtype=float)
        h_K_like = B_h.T.dot(m_C).dot(B_h)
        h_K_like += a_m.dot(a_m.T)
        h_h_like = -B_h.T.dot(m_C).dot(b_h)

        tmp = np.zeros((self.dim, 1), dtype=float)
        tmp[0, 0] = m_m[-1, 0] / (1 - np.sum(m_m[:-1, 0]))
        h_h_like += a_m.dot(a_m.T).dot(tmp)

        h_K_like /= variation
        h_h_like /= variation

        self.surfel.bel_h_K -= self.surfel.like_outmsg_h_K[i]
        self.surfel.bel_h_h -= self.surfel.like_outmsg_h_h[i]
        self.surfel.like_outmsg_h_K[i] = h_K_like
        self.surfel.like_outmsg_h_h[i] = h_h_like
        self.surfel.bel_h_K += self.surfel.like_outmsg_h_K[i]
        self.surfel.bel_h_h += self.surfel.like_outmsg_h_h[i]
        self.surfel.bel_h_C = np.linalg.inv(self.surfel.bel_h_K)
        self.surfel.bel_h_m = self.surfel.bel_h_C.dot(self.surfel.bel_h_h)

        # Update q(v) --------------------------------------------------
        if self.dim == 3:
            a_h = np.array(
                [
                    [self.surfel.bel_h_m[0, 0] - self.surfel.bel_h_m[1, 0]],
                    [self.surfel.bel_h_m[0, 0] - self.surfel.bel_h_m[2, 0]],
                    [1],
                ],
                dtype=float,
            )
        else:
            a_h = np.array(
                [[self.surfel.bel_h_m[0, 0] - self.surfel.bel_h_m[1, 0]], [1]], dtype=float
            )

        u_h = np.zeros((self.dim, 1), dtype=float)
        u_h[-1, 0] = self.surfel.bel_h_m[0, 0]
        A_h = a_h.dot(a_h.T)

        E = np.trace(np.linalg.multi_dot((B_m.T, self.surfel.bel_h_C, B_m, m_C)))
        E += np.linalg.multi_dot((m_m.T, B_m.T, self.surfel.bel_h_C, B_m, m_m))
        E += 2 * np.linalg.multi_dot((b_m.T, self.surfel.bel_h_C, B_m, m_m))
        E += np.linalg.multi_dot((b_m.T, self.surfel.bel_h_C, b_m))
        E += np.trace(A_h.dot(m_C))
        E += np.linalg.multi_dot((m_m.T, A_h, m_m))
        E += -2 * np.linalg.multi_dot((u_h.T, A_h, m_m))
        E += np.linalg.multi_dot((u_h.T, A_h, u_h))
        E *= 0.5

        self.surfel.bel_v_a -= self.surfel.like_outmsg_v_a[i]
        self.surfel.bel_v_b -= self.surfel.like_outmsg_v_b[i]

        self.surfel.like_outmsg_v_b[i] = E.item()
        self.surfel.like_outmsg_v_a[i] = 0.5

        self.surfel.bel_v_a += self.surfel.like_outmsg_v_a[i]
        self.surfel.bel_v_b += self.surfel.like_outmsg_v_b[i]


import csv
import math
import sys

import pss
import psqi

# indexes of field names
timestamp = 0
grade = 1
exercise = 31
screen_time = 32
extracurr = 33

# validate cmdline arguments
# i know it's not robust, but i probably shouldn't overengineer
# some code for an AP stats class which only I will ever run
if len(sys.argv) < 3:
    print("usage: $ python3 score.py input.csv output.csv [--no-print]")
    sys.exit()

printing = True
if len(sys.argv) == 4 and sys.argv[3] == '--no-print':
    printing = False

pss_scores = []
psqi_scores = []

# main routine
with open(sys.argv[1], "r", encoding='utf8') as fin:
    reader = csv.reader(fin)
    next(reader)

    with open(sys.argv[2], "w", encoding='utf8') as fout:
        writer = csv.writer(fout)

        # write output field names
        fout.write("Timestamp,Grade,PSS,PSQI,Exercise,Screen Time,Extracurriculars\n")

        print("**************************************")
        for response in reader:
            pss_score = pss.calc_pss(response)
            psqi_score = psqi.calc_psqi(response, printing)

            pss_scores.append(pss_score)
            psqi_scores.append(psqi_score)

            # if printing:
            print('| ' + response[timestamp], end='')
            print(f"  PSS {pss_score:02d}", end='')
            print(f"  PSQI {psqi_score:02d} |")
            print("**************************************")

            row = [
                response[timestamp], response[grade],
                pss_score, psqi_score, response[exercise],
                response[screen_time], response[extracurr]
                ]
            writer.writerow(row)

def mean(ls):
    total = 0
    for n in ls:
        total += n
    return total / len(ls)

def std_dev(ls, mean):
    total = 0
    for n in ls:
        total += ((n - mean) ** 2)
    var = total / (len(ls) - 1)
    return math.sqrt(var)

print(f"\n\nSCORING COMPLETE.")
print(f"{len(pss_scores):d} SUBJECTS ANALYZED.\n")

mean_psqi = mean(psqi_scores)
mean_pss = mean(pss_scores)

print(f"MEAN    PSQI  {mean_psqi:.02f}")
print(f"MEAN    PSS   {mean_pss:.02f}")
print(f"STD_DEV PSQI  {std_dev(psqi_scores, mean_psqi):.02f}")
print(f"STD_DEV PSS   {std_dev(pss_scores, mean_pss):.02f}")


import copy
import json
import os.path
import torch
from transformers import BartTokenizer
from typing import List, Dict
from fire import Fire
from transformers import AutoTokenizer, PreTrainedTokenizer

def fix_random(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.enabled = False

def load_json(file_path):
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            data.append(json.loads(line.strip()))
    return data

def write_json(file_path, data: List[Dict]):
    with open(file_path, 'w') as f:
        for line in data:
            f.write(json.dumps(line))
            f.write('\n')

def norm_strategy(strategy):
    norm_str = "-".join(strategy.split())
    return "@["+norm_str+"]"

def get_strategy(file_path, norm=False):
    with open(file_path,'r', encoding='utf-8') as f:
        data = json.load(f)
    data = [d.replace('[','').replace(']','') for d in data]
    if norm:
        data = [norm_strategy(d) for d in data]
    print('strategy: ', data)

    return data

def _norm(x):
    return ' '.join(x.strip().split()).lower()

#todo: add <helper> and <seeker> tokens to differentiate between the two (sometimes one speaker talks multiple times in the dataset)
def construct_conversational_dataset(
        file_path: str,
        tokenizer: BartTokenizer,
        add_cause: bool = False,
        with_strategy: bool = False,
        joint_strategy_utt: bool = True,
        load_from_cache: bool = True,
) -> str:

    split = 'train'
    if 'valid' in file_path:
        split = 'valid'
    elif 'test' in file_path:
        split = 'test'

    cached_preprocessed_file = file_path.replace(split, f"preprocessed_{split}")

    if os.path.exists(cached_preprocessed_file) and load_from_cache:
        print(f"loading preprocessed {split} dataset from {cached_preprocessed_file}")
        return cached_preprocessed_file

    data = load_json(file_path)
    print(f"parsing {file_path} file")
    sep_token = tokenizer.sep_token if tokenizer.sep_token is not None else " "
    print(f"setting {sep_token} as sep token")

    total_data = []
    is_train = 'train' in file_path

    # todo: check if required later:
    # if 'test' in file_path and self.model_type > 4:
    #     gt_strategy = read_pk('./final_data/test_extend_label.pk')
    #     if lookahead is not True:
    #         print("do not use lookahead! ")
    #         predict_strategy = read_pk('./final_data/wo_lookahead_predicted.pk')
    #     else:
    #         print('use lookahead! ')
    #         predict_strategy = read_pk('./final_data/multiesc_predicted_strategy.pk')
    #     print("acc: ", accuracy_score(gt_strategy, predict_strategy))

    predict_strategy_index = 0
    for case_example in data:
        dialog = case_example['dialog']
        dialog_len = len(dialog)
        emotion_type = case_example['emotion_type']
        problem_type = case_example['problem_type']
        situation = case_example['situation']
        tot_strategy = []
        for index, tmp_dic in enumerate(dialog):
            if tmp_dic['speaker'] == 'sys' and tmp_dic['strategy'] != "Others":
                # todo: doesnt it mess with data?
                tot_strategy.append(norm_strategy(tmp_dic['strategy']))

        # initial history is emotion_type + problem_type + situation
        history = [_norm(emotion_type) + sep_token + _norm(
            problem_type) + sep_token + _norm(situation)]

        tmp_strategy_list = []
        for index, tmp_dic in enumerate(dialog):
            text = _norm(tmp_dic['text'])
            if index == 0 and tmp_dic['speaker'] != 'sys':
                # todo: why prepend it to history?
                # handling when conversation starts with help seeker
                history[0] = "<seeker> " + text + sep_token + history[0]
                continue
            if tmp_dic['speaker'] == 'sys' and tmp_dic['strategy'] != "Others":
                # build one example from history up to this point of conversation
                tmp_strategy = norm_strategy(tmp_dic['strategy'])
                save_s = [x for x in tot_strategy[len(tmp_strategy_list):]].copy()
                assert len(save_s) > 0, print(tot_strategy, tmp_strategy_list)
                tmp_history = copy.deepcopy(history)
                if joint_strategy_utt:
                    response = tmp_strategy + " " + text
                else:
                    response = text

                # todo: check if required later:
                # if with_strategy and self.model_type > 4:
                #     if 'test' in file_path:
                #         # tmp_history[-1] = tmp_history[-1] + self.sep_token + predict_strategy[predict_strategy_index]
                #         if self.model_type == 8:
                #             # response = predict_strategy[predict_strategy_index] + " " + text
                #             tmp_history.append(predict_strategy[predict_strategy_index])
                #         else:
                #             tmp_history[-1] = tmp_history[-1] + self.sep_token + predict_strategy[
                #                 predict_strategy_index]
                #         predict_strategy_index += 1
                #     else:
                #         if self.model_type == 8:
                #             tmp_history.append(tmp_strategy)
                #             # response = tmp_strategy + " " + text
                #         else:
                #             tmp_history[-1] = tmp_history[-1] + self.sep_token + tmp_strategy

                total_data.append({
                    "history": tmp_history,
                    "strategy": tmp_strategy,
                    "history_strategy": tmp_strategy_list,
                    "response": response,
                    "future_strategy": ' '.join(save_s),
                    # todo: what is the use case for stage of conversation?
                    "stage": 5 * index // dialog_len,
                })
                tmp_strategy_list.append(tmp_strategy)

            if tmp_dic['speaker'] == 'sys':
                # handling helper utterance with strategy
                tmp_strategy = norm_strategy(tmp_dic['strategy'])
                if with_strategy:
                    # add strategy as control code to the history text
                    tmp_sen = "<helper> " + tmp_strategy + " " + text
                    history.append(tmp_sen)
                else:
                    history.append("<helper> " + text)
            else:
                # help seeker utterance
                utt = text

                if add_cause:
                    cause = tmp_dic['cause']
                    if cause is not None and len(cause) > 0:
                        utt = cause + sep_token + utt

                history.append("<seeker> " + utt)

    # random_idx = random.randint(0, len(total_data) - 1)
    # print(f'printing training example {random_idx}:')
    # print(total_data[random_idx])

    write_json(cached_preprocessed_file, total_data)
    return cached_preprocessed_file

class InputPreprocessor:
    def __init__(
            self,
            preprocessor_type: str,
            tokenizer: PreTrainedTokenizer = None,
            max_source_length: int = None,
            max_target_length: int = None,
    ):
        self.preprocessor_type = preprocessor_type
        self.tokenizer = tokenizer
        self.max_source_length = max_source_length
        self.max_target_length = max_target_length

        processing_functions = {
            'joint_strategy_utterance_generation': self.joint_strategy_utterance_generation,
            'strategy_generation': self.strategy_generation_tokenization,
        }

        self.preprocess = processing_functions[preprocessor_type]

    def strategy_generation_tokenization(self, example):
        history = example['history']
        target = example['future_strategy']
        full_text = self.tokenizer.sep_token.join(history)

        inputs = self.tokenizer(full_text, add_special_tokens=True, max_length=self.max_source_length, truncation=True)
        labels = self.tokenizer(target, add_special_tokens=True, max_length=self.max_target_length, truncation=True)

        return {
            'input_ids': inputs['input_ids'],
            'attention_mask': inputs['attention_mask'],
            'labels': labels['input_ids'],
        }

    def joint_strategy_utterance_generation(self, example):
        history = example['history']
        target = example['response']
        full_text = self.tokenizer.bos_token + " ".join(history) + " <helper> "

        inputs = self.tokenizer(full_text, add_special_tokens=False, max_length=self.max_source_length, truncation=True)
        labels = self.tokenizer(target, add_special_tokens=True, max_length=self.max_target_length, truncation=True)

        return {
            'input_ids': inputs['input_ids'],
            'attention_mask': inputs['attention_mask'],
            'labels': labels['input_ids'],
        }

# class CustomDataCollator(DataCollator):
#     def __init__(self, tokenizer):
#         self.tokenizer = tokenizer
#
#     def sequence_only_strategy_generation_collator(self, batch):
#         pass

def main(
    base_file_path: str,
    tokenizer_name_or_path: str = 'facebook/bart-base',
    add_cause: bool = False,
    with_strategy: bool = False,
):
    splits = ['train', 'valid', 'test']
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
    tokenizer.add_special_tokens({'sep_token': '<sep>'})
    for split in splits:
        file_path = os.path.join(base_file_path, f"{split}.json")
        construct_conversational_dataset(
            file_path,
            tokenizer,
            add_cause=add_cause,
            with_strategy=with_strategy,
            load_from_cache=False,
        )

if __name__ == "__main__":
    Fire(main)


import logging
logging.basicConfig(format="[%(asctime)s] %(message)s", datefmt="%m-%d %H:%M:%S")

import tensorflow as tf
import numpy as np

logger = logging.getLogger(__name__)

def get_shape(layer):
    return layer.get_shape().as_list()

def skew(inputs, scope="skew"):
    with tf.name_scope(scope):
        batch, height, width, channel = get_shape(inputs)
        rows = tf.split(inputs, height, 1)

        new_width = width + height - 1
        new_rows = []

        for idx, row in enumerate(rows):
            transposed_row = tf.transpose(tf.squeeze(row, [1]), [0, 2, 1]) # [batch, channel, width]
            squeezed_row = tf.reshape(transposed_row, [-1, width]) # [batch*channel, width]
            padded_row = tf.pad(squeezed_row, [[0, 0], [idx, height - 1 - idx]]) # [batch*channel, width]

            unsqueezed_row = tf.reshape(padded_row, [-1, channel, new_width]) # [batch, channel, new_width]
            untransposed_row = tf.transpose(unsqueezed_row, [0, 2, 1]) # [bacth, new_width, channel]

            assert get_shape(untransposed_row) == [batch, new_width, channel], "wrong shape of skewed row"
            new_rows.append(untransposed_row)

        outputs = tf.stack(new_rows, axis=1, name="output")
        assert get_shape(outputs) == [batch, height, new_width, channel], "wrong shape of outputs"

    logger.debug('[skew] %s : %s %s -> %s %s' \
        % (scope, inputs.name, inputs.get_shape(), outputs.name, outputs.get_shape()))
    return outputs

def unskew(inputs, width=None, scope="unskew"):
    with tf.name_scope(scope):
        batch, height, skewed_width, channel = get_shape(inputs)
        width = width if width else height

        new_rows = []
        rows = tf.split(inputs, height, 1)

        for idx, row in enumerate(rows):
            new_rows.append(tf.slice(row, [0, 0, idx, 0], [-1, -1, skewed_width - width + 1,-1]))
        # outputs = tf.stack(new_rows, axis=1, name="output")
        outputs = tf.concat(new_rows, axis=1, name="output")
    logger.debug('[unskew] %s : %s %s -> %s %s'\
        % (scope, inputs.name, inputs.get_shape(), outputs.name, outputs.get_shape()))
    return outputs

WEIGHT_INITIALIZER = tf.contrib.layers.xavier_initializer()

def conv2d(
    inputs,
    num_outputs,
    kernel_shape, # [kernel_height, kernel_width]
    mask_type, # None, "A" or "B",
    strides=[1, 1], # [column_wise_stride, row_wise_stride]
    padding="SAME",
    activation_fn=None,
    weights_initializer=WEIGHT_INITIALIZER,
    weights_regularizer=None,
    biases_initilizer=tf.zeros_initializer(),
    biases_regularizer=None,
    scope="conv2d"):
    with tf.variable_scope(scope):
        mask_type = mask_type.lower()
        batch, height, width, channel = inputs.get_shape().as_list()

        kernel_h, kernel_w = kernel_shape
        stride_h, stride_w = strides

        assert kernel_h % 2 == 1 and kernel_w % 2 == 1, \
        "kernel height and width should be odd number"

        center_h = kernel_h // 2
        center_w = kernel_w // 2

        weights_shape = [kernel_h, kernel_w, channel, num_outputs]
        weights = tf.get_variable("weights", weights_shape,
        tf.float32, weights_initializer, weights_regularizer)

        if mask_type is not None:
            mask = np.ones((kernel_h, kernel_w, channel, num_outputs), dtype=np.float32)
            mask[center_h, center_w+1:,:, :] = 0.
            mask[center_h+1, :, :, :] = 0.

            if mask_type == 'a':
                mask[center_h, center_w, :, :] = 0.

            weights *= tf.constant(mask, dtype=tf.float32)
            tf.add_to_collection('conv2d_weights_%s' % mask_type, weights)

        outputs = tf.nn.conv2d(inputs, weights, [1, stride_h, stride_w, 1], padding=padding, name='outputs')
        tf.add_to_collection('conv2d_outputs', outputs)

        if biases_initilizer != None:
            biases = tf.get_variable("biases", [num_outputs,],
                tf.float32, biases_initilizer, biases_regularizer)
            outputs = tf.nn.bias_add(outputs, biases, name='outputs_plus_b')

        if activation_fn:
            outputs = activation_fn(outputs, name='outputs_with_fn')

        logger.debug('[conv2d_%s] %s : %s %s -> %s %s' \
            % (mask_type, scope, inputs.name, inputs.get_shape(), outputs.name, outputs.get_shape()))

        return outputs


"""

"""

from skimpy.core.parameters import ParameterValuePopulation
from scipy.stats import multivariate_normal

import tensorflow as tf

import pandas as pd
import numpy as np

EPSILON = 1e-9

class SecureMultivariateNormal(object):
    def __init__(self, mu, sigma, var):
        self.variable_parameters = var > EPSILON
        self.constant_parameters = var < EPSILON

        self.mu = mu

        self.variable_index = var.index[self.variable_parameters]
        self.const_index = var.index[self.constant_parameters]

        # TODO: RAISE A MORE INFORMATIVE ERROR WHEN THE COV IS SIGULAR!

        self._dist = multivariate_normal(mu[self.variable_parameters],
                                         sigma.loc[self.variable_parameters,self.variable_parameters])

    def rvs(self,size,random_state=None):
        values_var = self._dist.rvs(size=size, random_state=random_state)
        if size > 1:
            df = pd.DataFrame(values_var, columns=self.variable_index)
        else:
            df = pd.DataFrame(values_var, index=self.variable_index, columns=[0]).T
        values_cons = pd.concat( [ self.mu[self.constant_parameters] , ]*size , axis=1)

        return pd.concat([df, values_cons.T], axis=1)

class LogNormalPriorParameterDistribution():
    """
    TF Based model
    """
    pass

class PosteriorLogNormalParameterPopulation(object):
    """

    """
    def __init__(self, parameter_poulations, likelyhoods=None):
        self.mu = []
        self.sigma = []
        self.pdf = []

        for pop in parameter_poulations:
            var = pop.log_var()
            mu = pop.log_mean()
            sigma = pop.log_cov()

            self.mu.append( mu )
            self.sigma.append( sigma )
            self.pdf.append( SecureMultivariateNormal(mu,sigma,var) )

        N = len(parameter_poulations)
        if likelyhoods is None:
            self.weights = np.ones(N)/N
        else:
            self.weights = likelyhoods/sum(likelyhoods)

        self.cum_weights = np.cumsum(self.weights)

    def resample(self, N, seed=None):
        if not seed is None:
            np.random.seed(seed=seed)

        # Gillespie type of algorithm for resampling
        x = []
        for s in range(N):
            r = np.random.rand()
            j = sum(r > self.cum_weights)
            x_i = np.exp( self.pdf[j].rvs(size=1, random_state=None) )
            x.append(x_i)
        df = pd.concat(x, axis=0, ignore_index=True, )
        return df

class PosteriorNormalParameterPopulation(object):
    """

    """
    def __init__(self, parameter_poulations, likelyhoods=None):
        self.mu = []
        self.sigma = []
        self.pdf = []

        for pop in parameter_poulations:
            var = pop.var()
            mu = pop.mean()
            sigma = pop.cov()

            self.mu.append( mu )
            self.sigma.append( sigma )
            self.pdf.append( SecureMultivariateNormal(mu,sigma,var) )

        N = len(parameter_poulations)
        if likelyhoods is None:
            self.weights = np.ones(N)/N
        else:
            self.weights = likelyhoods/sum(likelyhoods)

        self.cum_weights = np.cumsum(self.weights)

    def resample(self, N, seed=None):
        if not seed is None:
            np.random.seed(seed=seed)

        # Gillespie type of algorithm for resampling
        x = []
        for s in range(N):
            r = np.random.rand()
            j = sum(r > self.cum_weights)
            x_i = self.pdf[j].rvs(size=1, random_state=None)
            x.append(x_i)
        df = pd.concat(x, axis=0, ignore_index=True, )
        return df


import sys
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QLineEdit, QLabel, QGraphicsScene, QGraphicsView, \
    QGraphicsItem, QMainWindow, QFileDialog, QVBoxLayout
from PyQt5.QtGui import QIcon, QPixmap, QPainter, QBrush, QPen, QMouseEvent, QImage, QKeySequence, QCursor
from PyQt5.QtCore import pyqtSlot, Qt, QDir, QRect
from PyQt5 import QtWidgets

#from printscreen import PrintScreen

#import keyboard

class App(QMainWindow):

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

        self.title = 'Screen Shot'

        self.initUI()

    def initUI(self):

        self.setWindowTitle(self.title)

        self.setMinimumSize(50, 50)

        self.show()

def open_print_screen():
    print("open_print")
    img = QApplication.primaryScreen().grabWindow(0)
    PrintScreen(img)

if __name__ == '__main__':

    app = QApplication(sys.argv)
    ex = App()
    sys.exit(app.exec_())


from __future__ import print_function

import chess
import chess.uci
import chess.variant
import time
import textwrap
import argparse
import itertools
import logging
import sys

def test_epd(engine, epd, VariantBoard, movetime):
    position = VariantBoard()
    epd_info = position.set_epd(epd)
    epd_string = "%s" % epd_info.get("id", position.fen())
    if "am" in epd_info:
        epd_string = "%s (avoid %s)" % (epd_string, " and ".join(position.san(am) for am in epd_info["am"]))
    if "bm" in epd_info:
        epd_string = "%s (expect %s)" % (epd_string, " or ".join(position.san(bm) for bm in epd_info["bm"]))

    engine.ucinewgame()
    engine.setoption({"UCI_Variant": VariantBoard.uci_variant})
    engine.position(position)

    enginemove, pondermove = engine.go(movetime=movetime)

    if "am" in epd_info and enginemove in epd_info["am"]:
        print("%s: %s | +0" % (epd_string, position.san(enginemove)))
        return 0.0
    elif "bm" in epd_info and not enginemove in epd_info["bm"]:
        print("%s: %s | +0" % (epd_string, position.san(enginemove)))
        return 0.0
    else:
        print("%s: %s | +1" % (epd_string, position.san(enginemove)))
        return 1.0

def test_epd_with_fractional_scores(engine, epd, VariantBoard, movetime):
    info_handler = chess.uci.InfoHandler()
    engine.info_handlers.append(info_handler)

    position = VariantBoard()
    epd_info = position.set_epd(epd)
    epd_string = "%s" % epd_info.get("id", position.fen())
    if "am" in epd_info:
        epd_string = "%s (avoid %s)" % (epd_string, " and ".join(position.san(am) for am in epd_info["am"]))
    if "bm" in epd_info:
        epd_string = "%s (expect %s)" % (epd_string, " or ".join(position.san(bm) for bm in epd_info["bm"]))

    engine.ucinewgame()
    engine.setoption({"UCI_Variant": VariantBoard.uci_variant})
    engine.position(position)

    # Search in background
    search = engine.go(infinite=True, async_callback=True)

    score = 0.0

    print("%s:" % epd_string, end=" ")
    sys.stdout.flush()

    for step in range(0, 3):
        time.sleep(movetime / 4000.0)

        # Assess the current principal variation.
        with info_handler as info:
            if 1 in info["pv"] and len(info["pv"][1]) >= 1:
                move = info["pv"][1][0]
                print("(%s)" % position.san(move), end=" ")
                sys.stdout.flush()
                if "am" in epd_info and move in epd_info["am"]:
                    continue #fail
                elif "bm" in epd_info and not move in epd_info["bm"]:
                    continue #fail
                else:
                     score = 1.0 / (4 - step)
            else:
                print("(no pv)", end=" ")
                sys.stdout.flush()

    # Assess the final best move by the engine.
    time.sleep(movetime / 4000.0)
    engine.stop()
    enginemove, pondermove = search.result()
    if "am" in epd_info and enginemove in epd_info["am"]:
        pass #fail
    elif "bm" in epd_info and not enginemove in epd_info["bm"]:
        pass #fail
    else:
         score = 1.0

    print("%s | +%g" % (position.san(enginemove), score))

    engine.info_handlers.remove(info_handler)
    return score

if __name__ == "__main__":
    # Parse command line arguments.
    parser = argparse.ArgumentParser(description="Run an EPD test suite with an UCI engine.")
    parser.add_argument("-e", "--engine", required=True,
        help="The UCI engine under test.")
    parser.add_argument("epd", nargs="+", type=argparse.FileType("r"),
        help="EPD test suite(s).")
    parser.add_argument("-v", "--variant", default="standard",
        help="Use a non-standard chess variant.")
    parser.add_argument("-t", "--movetime", default=1000, type=int,
        help="Time to move in milliseconds.")
    parser.add_argument("-s", "--simple", dest="test_epd", action="store_const",
        default=test_epd_with_fractional_scores,
        const=test_epd,
        help="Run in simple mode without fractional scores.")
    parser.add_argument("-d", "--debug", action="store_true",
        help="Show debug logs.")
    args = parser.parse_args()

    # Configure logger.
    logging.basicConfig(level=logging.DEBUG if args.debug else logging.WARNING)

    # Find variant.
    VariantBoard = chess.variant.find_variant(args.variant)

    # Open engine.
    engine = chess.uci.popen_engine(args.engine)
    engine.uci()

    # Run each test line.
    score = 0.0
    count = 0

    for epd in itertools.chain(*args.epd):
        # Skip comments and empty lines.
        epd = epd.strip()
        if not epd or epd.startswith("#") or epd.startswith("%"):
            print(epd.rstrip())
            continue

        # Run the actual test.
        score += args.test_epd(engine, epd, VariantBoard, args.movetime)
        count += 1

    engine.quit()

    print("-------------------------------")
    print("%g / %d" % (score, count))


import json

from .icloud_sample import SampleICloudApplication

from .. import utils

class SampleLiveICloudApplication(SampleICloudApplication):
    display_name = 'Sample Live Data Application'

    def run(self):
        # Register the account for the iCloud service.
        self.client.register_account(self.account)
        # Attempt to login to the account.
        self.log_in()
        # Choose a data type and retrieve it.
        self.fetch_data()
        # Wait for any pending tasks to complete
        self.client.wait_for_pending_tasks()
        utils.info_message('All tasks completed')

    def fetch_data(self):
        """Prompt for a data type choice and execute the `fetch_data` task.
        The results are saved to a file in json format.
        """
        choices = self.available_data
        choices.insert(0, 'All')

        selected_data_type = utils.select_item(
            choices,
            'Please select what data to fetch:',
            'Available data:',
        )

        if selected_data_type == 'All':
            selected_data_type = ','.join(self.available_data)

        utils.pending_message('Performing fetch data task...')

        fetch_data_task = self.client.data(
            account=self.account,
            data=selected_data_type,
        )

        # Wait here for result as rest of sample app relies on it.
        fetch_data_task.wait_for_result(timeout=self.timeout)
        fetch_data_result = json.loads(fetch_data_task.result)

        # Write the result to file.
        task_id = fetch_data_task.uuid
        filepath = utils.get_or_create_filepath('%s.json' % task_id)

        with open(filepath, 'w') as out:
            json.dump(fetch_data_result, out, indent=2)

        utils.info_message('Fetch data successful. Output file: %s.json' % task_id)

        return fetch_data_result

    @property
    def available_data(self):
        live_feeds = [
            'location',
            'mobileme_contacts',
            'web_browser_history',
            'live_call_history',
            'live_photos',
        ]
        return [c for c in self.client.available_data
                if c in live_feeds]


import math,sys,uuid
import numpy as np
import numpy.matlib
from multiprocessing import Pool
from multiprocessing import shared_memory
from calibration_tools import *

def globalize(func):
  def result(*args, **kwargs):
    return func(*args, **kwargs)
  result.__name__ = result.__qualname__ = uuid.uuid4().hex
  setattr(sys.modules[result.__module__], result.__name__, result)
  return result

def analysis_uvwdir_loop(skymodel,clusterfile,uvwfile,rhofile,solutionsfile,z_solfile,flow=110,fhigh=170,ra0=0,dec0=math.pi/2,tslots=10,alpha=0.1,Nparallel=4):
    # ra0,dec0: phase center (rad)
    # tslots: -t option
    # alpha: spatial constraint regularization parameter
    # Nparallel=number of parallel jobs to use
    flow=flow*1e6
    fhigh=fhigh*1e6

    # if 1, IQUV, else only I
    fullpol=0
    loop_in_r=False # use 8 blocks instead of looping

    # read Z (global) solutions to get metadata for consensus polynomial
    N,f0,Ne,K1,Z1=read_global_solutions(z_solfile)
    # N: stations
    # Ne: consensus poly terms, same as -P parameter in sagecal
    # baselines
    B=int(N*(N-1)/2)
    # reference freq (for consensus polynomial)
    assert(f0>=flow and f0<=fhigh) # mean of all freqs
    #%%%%%%%%%%%%%%%%%% consensus polynomial info
    Nf=8 # no. of freqs: make sure to match all data
    f=np.linspace(flow,fhigh,Nf)
    polytype=1 # 0: ordinary, 1: Bernstein

    # read solutions file (also get the frequency(MHz)) J: Kx2N Nt x 2 (2Nx2 blocks Nt times)
    freq,J=readsolutions(solutionsfile)
    # read sky model Ct: Kx T x 4 (each row XX,XY,YX,YY)
    K,Ct=skytocoherencies(skymodel,clusterfile,uvwfile,N,freq,ra0,dec0)
    assert(K==K1)

    # ADMM rho, per each direction, scale later
    # scale rho linearly with sI
    rho=read_rho(rhofile,K)

    # read u,v,w,xx(re,im), xy(re,im) yx(re,im) yy(re,im)
    XX,XY,YX,YY=readuvw(uvwfile)
    # create shared memory equal to XX,XY,YX,YY buffers for parallel processing
    shmXX=shared_memory.SharedMemory(create=True,size=XX.nbytes)
    shmXY=shared_memory.SharedMemory(create=True,size=XY.nbytes)
    shmYX=shared_memory.SharedMemory(create=True,size=YX.nbytes)
    shmYY=shared_memory.SharedMemory(create=True,size=YY.nbytes)
    # create arrays that can be used in multiprocessing
    XX0=np.ndarray(XX.shape,dtype=XX.dtype,buffer=shmXX.buf)
    XY0=np.ndarray(XY.shape,dtype=XY.dtype,buffer=shmXY.buf)
    YX0=np.ndarray(YX.shape,dtype=YX.dtype,buffer=shmYX.buf)
    YY0=np.ndarray(YY.shape,dtype=YY.dtype,buffer=shmYY.buf)
    # how many timeslots to use per calibration (-t option)
    T=tslots
    Ts=int(XX.shape[0]//(B*T))

    # check this agrees with solutions
    nx,ny=J[0].shape
    if nx<2*N*Ts:
     print('Error: solutions size does not match with data size')
     exit

    # which frequency index to work with
    fidx=np.argmin(np.abs(f-freq))

    # addition to Hessian
    Hadd=np.zeros((K,4*N,4*N),dtype=np.float32)
    for ci in range(K):
     # note: F is dependent on rho when alpha!=0
     # example: making F=rand(2N,2N) makes performance worse
     F,P=consensus_poly(Ne,N,f,f0,fidx,polytype=polytype,rho=rho[ci],alpha=alpha)
     FF=np.matmul(F.transpose(),F)
     if alpha>0.0:
       PP=np.matmul(P.transpose(),P)
       H11=0.5*rho[ci]*FF+0.5*alpha*rho[ci]*rho[ci]*PP
       H12=0.5*FF+0.5*alpha*rho[ci]*PP
       H21=H12
       H22=-0.5/rho[ci]*(np.eye(2*N)-FF)+0.5*alpha*PP
       Htilde=H11-np.matmul(H12,np.matmul(np.linalg.pinv(H22),H21))
       Hadd[ci]=np.kron(np.eye(2),Htilde)
     else:
       Hadd[ci]=0.5*rho[ci]*np.kron(np.eye(2),np.matmul(FF,np.eye(2*N)+np.matmul(np.linalg.pinv(np.eye(2*N)-FF),FF)))

############################# loop over timeslots
############################# local function
    @globalize
    def process_chunk(ncal):
        ts=ncal*T
        print('%d %d %d'%(ts,Ts,ncal))
        R=np.zeros((2*B*T,2),dtype=np.csingle)
        R[0:2*B*T:2,0]=XX[ts*B:ts*B+B*T]
        R[0:2*B*T:2,1]=XY[ts*B:ts*B+B*T]
        R[1:2*B*T:2,0]=YX[ts*B:ts*B+B*T]
        R[1:2*B*T:2,1]=YY[ts*B:ts*B+B*T]

        # D_Jgrad K x 4Nx4N tensor
        H=Hessianres(R,Ct[:,ts*B:ts*B+B*T],J[:,ncal*2*N:ncal*2*N+2*N],N)
        H+=Hadd

        # set to zero
        XX0[ts*B:ts*B+B*T]=0
        XY0[ts*B:ts*B+B*T]=0
        YX0[ts*B:ts*B+B*T]=0
        YY0[ts*B:ts*B+B*T]=0

        if loop_in_r:
          for r in range(8):
            # dJ: K x 4NxB tensor
            dJ=Dsolutions(Ct[:,ts*B:ts*B+B*T],J[:,ncal*2*N:ncal*2*N+2*N],N,H,r)
            # dR: 4B x B (sum up all K)
            dR=Dresiduals(Ct[:,ts*B:ts*B+B*T],J[:,ncal*2*N:ncal*2*N+2*N],N,dJ,0,r) # 0 for not adding I to dR
            # find mean value over columns
            dR11=np.mean(dR[0:4*B:4],axis=0)
            dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
            XX0[ts*B:ts*B+B*T] +=dR11
            dR11=np.mean(dR[3:4*B:4],axis=0)
            dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
            YY0[ts*B:ts*B+B*T] +=dR11
            if fullpol:
              dR11=np.mean(dR[1:4*B:4],axis=0)
              dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
              XY0[ts*B:ts*B+B*T] +=dR11
              dR11=np.mean(dR[2:4*B:4],axis=0)
              dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
              YX0[ts*B:ts*B+B*T] +=dR11
        else:
          # dJ: 8 x K x 4NxB tensor
          dJ=Dsolutions_r(Ct[:,ts*B:ts*B+B*T],J[:,ncal*2*N:ncal*2*N+2*N],N,H)
          # dR: 8 x 4B x B (sum up all K)
          dR=Dresiduals_r(Ct[:,ts*B:ts*B+B*T],J[:,ncal*2*N:ncal*2*N+2*N],N,dJ,0) # 0 for not adding I to dR
          # find mean value over columns
          for r in range(8):
            dR11=np.mean(dR[r,0:4*B:4],axis=0)
            dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
            XX0[ts*B:ts*B+B*T] +=dR11
            dR11=np.mean(dR[r,3:4*B:4],axis=0)
            dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
            YY0[ts*B:ts*B+B*T] +=dR11
            if fullpol:
              dR11=np.mean(dR[r,1:4*B:4],axis=0)
              dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
              XY0[ts*B:ts*B+B*T] +=dR11
              dR11=np.mean(dR[r,2:4*B:4],axis=0)
              dR11=np.squeeze(np.matlib.repmat(dR11,1,T))
              YX0[ts*B:ts*B+B*T] +=dR11
############################# end local function
############################# loop over timeslots

    # create pool
    pool=Pool(Nparallel)
    pool.map(process_chunk,range(Ts))
    pool.close()
    pool.join()

    # copy back from shared memory
    XX[:]=XX0[:]
    XY[:]=XY0[:]
    YX[:]=YX0[:]
    YY[:]=YY0[:]
    # release shared memory
    shmXX.close()
    shmXX.unlink()
    shmXY.close()
    shmXY.unlink()
    shmYX.close()
    shmYX.unlink()
    shmYY.close()
    shmYY.unlink()

    scalefactor=8*(N*(N-1)/2)*T
    # scale by 8*(N*(N-1)/2)*T
    writeuvw('fff',scalefactor*XX,scalefactor*XY,scalefactor*YX,scalefactor*YY)

if __name__ == '__main__':
  # args skymodel clusterfile uvwfile rhofile solutionsfile z_solutions_file freq_low(MHz) freq_high(MHz) ra0 dec0 tslots alpha parallel_jobs
  import sys
  argc=len(sys.argv)
  if argc==13:
   analysis_uvwdir_loop(sys.argv[1],sys.argv[2],sys.argv[3],sys.argv[4],sys.argv[5],sys.argv[6],float(sys.argv[7]),float(sys.argv[8]),float(sys.argv[9]),float(sys.argv[10]),int(sys.argv[11]),float(sys.argv[12]))
  elif argc==14:
   analysis_uvwdir_loop(sys.argv[1],sys.argv[2],sys.argv[3],sys.argv[4],sys.argv[5],sys.argv[6],float(sys.argv[7]),float(sys.argv[8]),float(sys.argv[9]),float(sys.argv[10]),int(sys.argv[11]),float(sys.argv[12]),int(sys.argv[13]))
  else:
   print("Usage: python %s skymodel clusterfile uvwfile rhofile solutionsfile z_solutions_file freq_low freq_high ra0 dec0 tslots alpha parallel_jobs"%(sys.argv[0]))
  exit()