import csv
import json
import flask
from flask import Flask, request, jsonify
from flask_restful import Resource, Api
from __init__ import app, db
import math

from Core.airport import Airport
from Core.time import Time
from Core.carrier import Carrier
from Core.flights import Flights
from Core.delays_minutes import Delays_minutes
from Core.delays_amount import Delays_amount
from Core.delays import Delays
from Core.relation_table import Relation_table
from Core.statistics import Statistics

class CSVHandler():

    def getAirportCSV(dictionary = None, dataList = None):
        returnString = ""

        if not dictionary is None:

            returnString = returnString + dictionary["name"] + "\n"
            for uri in dictionary["uri-list"]:
                returnString = returnString + uri + "\n"
        else:
            for data in dataList:
                returnString = returnString + data["name"] + "," + data["uri"] + "\n"

        return returnString

    def getCarrierCSV(dictionary = None, dataList = None):
        returnString = ""

        if not dictionary is None:
            returnString = returnString + dictionary["name"] + "\n"
            returnString = returnString + dictionary["statistics-uri"] + "\n"
            for uri in dictionary["airport-uris"]:
                returnString = returnString + uri + "\n"

        else:
            for data in dataList:
                returnString = returnString + data["name"] + "," + data["uri"] + "\n"
        return returnString

    def getStatisticsCSV(dictionary):
        returnString = dictionary["flights-uri"] + ", " + dictionary["minutes-uri"] + ", " + dictionary["amount-uri"]
        return returnString

    def getMinutesCSV(dictionary):
        returnString = dictionary["month"] + ","
        returnString = returnString + str(dictionary["minutes-data"]["late-aircraft"]) + "," +  str(dictionary["minutes-data"]["carrier"]) + "," +  str(dictionary["minutes-data"]["security"]) + "," +  str(dictionary["minutes-data"]["weather"]) + "," +  str(dictionary["minutes-data"]["nas"]) + "," +  str(dictionary["minutes-data"]["total"])
        returnString = returnString + "\n" + dictionary["carrier-uri"]
        return returnString

    def getAmountCSV(dictionary):
        returnString = dictionary["month"] + ","
        returnString = returnString + str(dictionary["amount-data"]["late-aircraft"]) + "," +  str(dictionary["amount-data"]["carrier"]) + "," +  str(dictionary["amount-data"]["security"]) + "," +  str(dictionary["amount-data"]["weather"]) + "," +  str(dictionary["amount-data"]["nas"])
        returnString = returnString + "\n" + dictionary["carrier-uri"]
        return returnString

    def getFlightsCSV(dictionary):
        returnString = dictionary["month"] + ","
        returnString = returnString + str(dictionary["flights-data"]["cancelled"]) + "," +  str(dictionary["flights-data"]["ontime"]) + "," +  str(dictionary["flights-data"]["delayed"]) + "," +  str(dictionary["flights-data"]["diverted"]) + "," +  str(dictionary["flights-data"]["total"])
        returnString = returnString + "\n" + dictionary["carrier-uri"]
        return returnString


from torch import nn
import torch

class BaseConv(nn.Module):

    def __init__(self, inc, ouc, k, s):
        super().__init__()
        self.conv = nn.Conv2d(inc, ouc, k, s, (k-1)//2, 1, 1, False)
        self.norm = nn.BatchNorm2d(ouc)
        self.act = nn.ReLU(inplace=True)

    def forward(self, x):
        return self.act(self.norm(self.conv(x)))

class Bottleneck(nn.Module):

    def __init__(self, inc, ouc,
                 shortcut=True,
                 expansion=0.5,
                 depthwise=False):
        super().__init__()
        hidc = int(inc * expansion)
        # Conv = DWConv if depthwise else BaseConv
        Conv = BaseConv
        self.conv1 = BaseConv(inc, ouc, 1, s=1)
        self.conv2 = Conv(hidc, ouc, 3, s=1)
        self.use_add = shortcut and inc == ouc

    def forward(self, x):
        y = self.conv2(self.conv1(x))
        if self.use_add:
            y = y + x
        return y

class CSPLayer(nn.Module):

    def __init__(self, inc, ouc, depth=1, shortcut=True,
                 expand=0.5, depthwise=False):
        super().__init__()
        hidc = int(ouc * expand)
        self.conv1 = BaseConv(inc, hidc, 1, 1)
        self.conv2 = BaseConv(inc, hidc, 1, 1)
        self.conv3 = BaseConv(2 * hidc, ouc, 1, 1)
        module_list = [
            Bottleneck(hidc, hidc, shortcut, 1.0,
                       depthwise) for _ in range(depth)
        ]
        self.m = nn.Sequential(*module_list)

    def forward(self, x):
        x_1 = self.conv1(x)
        x_2 = self.conv2(x)
        x_1 = self.m(x_1)
        x = torch.cat((x_1, x_2), dim=1)
        return self.conv3(x)


#!/usr/bin/env python3

from relation import *
import unittest

class TestColumn(unittest.TestCase):
	def test_should_have_expected_defaults(self):
		col = Column('Name', str)
		self.assertEqual(col.name, 'Name')
		self.assertEqual(col.type, str)
		self.assertTrue(col.nullable)
		self.assertIsNone(col.index)

	def test_transform_should_replace_fields(self):
		col1 = Column('A', str, nullable=True, index=2)
		col2 = col1.transform(new_name='B', new_column_type=int,
				new_nullability=False, new_index=1)
		self.assertEqual(col2.name, 'B')
		self.assertEqual(col2.type, int)
		self.assertFalse(col2.nullable)
		self.assertEqual(col2.index, 1)

	def test_transform_should_not_modify_original_column(self):
		col1 = Column('A', str, nullable=True, index=2)
		col2 = col1.transform(new_name='B', new_column_type=int,
				new_nullability=False, new_index=1)
		self.assertEqual(col1.name, 'A')
		self.assertEqual(col1.type, str)
		self.assertTrue(col1.nullable)
		self.assertEqual(col1.index, 2)

	def test_transform_should_not_modify_unspecified_fields(self):
		col1 = Column('A', str, nullable=True, index=2)
		col2 = col1.transform()
		self.assertEqual(col1.name, 'A')
		self.assertEqual(col2.type, str)
		self.assertTrue(col2.nullable)
		self.assertEqual(col2.index, 2)

	def test_transform_should_set_zero_values(self):
		col1 = Column('A', str, nullable=True, index=2)
		col2 = col1.transform(new_nullability=False, new_index=0)
		self.assertFalse(col2.nullable)
		self.assertEqual(col2.index, 0)

class TestMaterialRelation(unittest.TestCase):

	def test_should_not_insert_wrong_number_of_attributes(self):
		relation = MaterialRelation([Column('uid', int), Column('name', str)])
		with self.assertRaisesRegex(TypeError, 'number of columns'):
			relation.insert([1, 'Alice', 3])

	def test_should_not_insert_value_with_wrong_type(self):
		relation = MaterialRelation([Column('id', int)])
		with self.assertRaisesRegex(TypeError, 'wrong type'):
			relation.insert(['hi'])

	def test_should_not_insert_null_into_not_null_column(self):
		relation = MaterialRelation([Column('id', int, nullable=False)])
		with self.assertRaisesRegex(TypeError, 'NULL'):
			relation.insert([None])

	def test_should_insert_null_into_nullable_column(self):
		relation = MaterialRelation([Column('id', int, nullable=True)])
		relation.insert([None])
		self.assertEqual([(None,)], list(relation))

	def test_should_have_expected_values_after_insert(self):
		relation = MaterialRelation([Column('uid', int), Column('name', str)])
		relation.insert((1, 'Alice'))
		relation.insert((2, 'Bob'))
		relation.insert((3, 'Eve'))

		self.assertEqual([(1, 'Alice'), (2, 'Bob'), (3, 'Eve')], list(relation))

	def test_should_have_expected_name(self):
		relation1 = MaterialRelation([Column('uid', int)], name='Users')
		relation2 = MaterialRelation([Column('uid', int)])
		self.assertEqual(relation1.name, 'Users')
		self.assertEqual(relation2.name, None)
		relation2.set_name('Users')
		self.assertEqual(relation2.name, 'Users')

class ValueExpression(Expression):
	def __init__(self, value, expression_type, nullable=True):
		self.value = value
		self.type = expression_type
		self.is_nullable = nullable

	def value_type(self):
		return self.type

	def nullable(self):
		return self.is_nullable

	def evaluate(self, row):
		return self.value

class TestConstantExpression(unittest.TestCase):
	def test_should_have_correct_attributes_for_value(self):
		expr = Constant(123)
		self.assertEqual(expr.value_type(), int)
		self.assertFalse(expr.nullable())
		self.assertEqual(expr.evaluate(('a', 2)), 123)

	def test_should_have_correct_attributes_for_null_value(self):
		expr = Constant(None)
		self.assertEqual(expr.value_type(), type(None))
		self.assertTrue(expr.nullable())
		self.assertEqual(expr.evaluate(('a', 2)), None)

class TestAttributeExpression(unittest.TestCase):
	def test_attributes_should_match_column(self):
		expr = Attribute(Column('uid', int, nullable=True, index=3))
		self.assertEqual(expr.value_type(), int)
		self.assertTrue(expr.nullable())

	def test_evaluate_should_return_error_for_incorrect_row_type(self):
		expr = Attribute(Column('uid', int, index=1))
		with self.assertRaisesRegex(TypeError, 'wrong type for column'):
			expr.evaluate((123, 'hello'))

	def test_should_evaluate_to_expected_value(self):
		expr = Attribute(Column('uid', int, nullable=True, index=1))
		self.assertEqual(expr.evaluate(('Alice', 123)), 123)
		self.assertEqual(expr.evaluate(('Bob', None)), None)

class TestComparison(unittest.TestCase):
	def test_should_non_null_values_should_produce_bool(self):
		cases = [
			('<', 1, 2, True),
			('<', 2, 2, False),
			('<', 2, 1, False),

			('<=', 1, 2, True),
			('<=', 2, 2, True),
			('<=', 2, 1, False),

			('=', 1, 2, False),
			('=', 2, 2, True),
			('=', 2, 1, False),

			('>=', 1, 2, False),
			('>=', 2, 2, True),
			('>=', 2, 1, True),

			('>', 1, 2, False),
			('>', 2, 2, False),
			('>', 2, 1, True),

			('<>', 1, 2, True),
			('<>', 2, 2, False),
			('<>', 2, 1, True),

			('!=', 1, 2, True),
			('!=', 2, 2, False),
			('!=', 2, 1, True),

			('<', 1.1, 2.2, True),
			('<', 2.2, 2.2, False),
			('<', 2.2, 1.1, False),

			('<=', 1.1, 2.2, True),
			('<=', 2.2, 2.2, True),
			('<=', 2.2, 1.1, False),

			('=', 1.1, 2.2, False),
			('=', 2.2, 2.2, True),
			('=', 2.2, 1.1, False),

			('>=', 1.1, 2.2, False),
			('>=', 2.2, 2.2, True),
			('>=', 2.2, 1.1, True),

			('>', 1.1, 2.2, False),
			('>', 2.2, 2.2, False),
			('>', 2.2, 1.1, True),

			('<>', 1.1, 2.2, True),
			('<>', 2.2, 2.2, False),
			('<>', 2.2, 1.1, True),

			('!=', 1.1, 2.2, True),
			('!=', 2.2, 2.2, False),
			('!=', 2.2, 1.1, True),

			('<', False, True, True),
			('<', True, True, False),
			('<', True, False, False),

			('<=', False, True, True),
			('<=', True, True, True),
			('<=', True, False, False),

			('=', False, True, False),
			('=', True, True, True),
			('=', True, False, False),

			('>=', False, True, False),
			('>=', True, True, True),
			('>=', True, False, True),

			('>', False, True, False),
			('>', True, True, False),
			('>', True, False, True),

			('<>', False, True, True),
			('<>', True, True, False),
			('<>', True, False, True),

			('!=', False, True, True),
			('!=', True, True, False),
			('!=', True, False, True),

			('<', 'Alice', 'Bob', True),
			('<', 'Bob', 'Bob', False),
			('<', 'Bob', 'Alice', False),

			('<=', 'Alice', 'Bob', True),
			('<=', 'Bob', 'Bob', True),
			('<=', 'Bob', 'Alice', False),

			('=', 'Alice', 'Bob', False),
			('=', 'Bob', 'Bob', True),
			('=', 'Bob', 'Alice', False),

			('>=', 'Alice', 'Bob', False),
			('>=', 'Bob', 'Bob', True),
			('>=', 'Bob', 'Alice', True),

			('>', 'Alice', 'Bob', False),
			('>', 'Bob', 'Bob', False),
			('>', 'Bob', 'Alice', True),

			('<>', 'Alice', 'Bob', True),
			('<>', 'Bob', 'Bob', False),
			('<>', 'Bob', 'Alice', True),

			('!=', 'Alice', 'Bob', True),
			('!=', 'Bob', 'Bob', False),
			('!=', 'Bob', 'Alice', True),
		]
		for op, lhs, rhs, expected_result in cases:
			expr = Comparison(op, Constant(lhs), Constant(rhs))
			self.assertEqual(expr.value_type(), bool)
			description = '%r %s %r' % (lhs, op, rhs)
			result = expr.evaluate([])
			self.assertIsNotNone(result, msg=description)
			self.assertEqual(result, expected_result, msg=description)
		# nullability, mixed types

	def test_should_return_null_if_operand_is_null(self):
		cases = [
			('<', None, 123, int),
			('=', 123, None, int),
			('<', None, 12.3, float),
			('<', None, 'hello', str),
			('=', None, False, bool),
			('<>', None, None, int),
			('!=', None, None, int),
			# Note: NULL <> NULL
			('=', None, None, int),
		]
		for op, lhs, rhs, operand_type in cases:
			expr = Comparison(op,
					ValueExpression(lhs, operand_type, nullable=True),
					ValueExpression(rhs, operand_type, nullable=True))
			self.assertTrue(expr.nullable())
			description = '%r %s %r %r' % (lhs, op, rhs, operand_type)
			result = expr.evaluate([lhs, rhs])
			self.assertIsNone(result, msg=description)

	def test_should_have_correct_nullability(self):
		cases = [
			(False, False, False),
			(False, True, True),
			(True, False, True),
			(True, True, True)
		]
		for lhs_null, rhs_null, result_null in cases:
			expr = Comparison('=',
					ValueExpression(0, int, nullable=lhs_null),
					ValueExpression(0, int, nullable=rhs_null))
			self.assertEqual(expr.nullable(), result_null)

	def test_should_return_error_if_operands_have_different_types(self):
		with self.assertRaisesRegex(TypeError, 'same type'):
			Comparison('=',
				ValueExpression(None, bool),
				ValueExpression(None, int))

class TestCastExpression(unittest.TestCase):
	def test_should_cast_bool_to_bool(self):
		for b in (False, True, None):
			expr = Cast(ValueExpression(b, bool, nullable=True), bool)
			self.assertEqual(expr.evaluate([]), b)

	def test_should_cast_integer_to_bool(self):
		cases = [
			(0, False),
			(1, True),
			(-123, True),
			(123, True),
			(None, None),
		]
		for i, b in cases:
			expr = Cast(ValueExpression(i, int, nullable=True), bool)
			self.assertEqual(expr.evaluate([]), b)

	def test_should_return_error_casting_float_to_bool(self):
		for f in (0.0, 1.0, -123.123, 3.14, None):
			with self.assertRaisesRegex(TypeError, 'FLOAT'):
				Cast(ValueExpression(f, float, nullable=True), bool)

	def test_should_cast_valid_string_to_bool(self):
		cases = [
			('false', False),
			('FALSE', False),
			('0', False),
			('true', True),
			('True', True),
			('1', True),
			(None, None)
		]
		for s, b in cases:
			expr = Cast(ValueExpression(s, str, nullable=True), bool)
			self.assertEqual(expr.evaluate([]), b)

	def test_should_return_error_casting_invalid_string_to_bool(self):
		for s in ('asf', '123', '1.0', 'unknown', 'null', ''):
			expr = Cast(ValueExpression(s, str, nullable=True), bool)
			with self.assertRaisesRegex(TypeError, 'invalid boolean'):
				expr.evaluate([])

	def test_should_cast_to_integer(self):
		cases = [
			(bool, False, 0),
			(bool, True, 1),
			(bool, None, None),
			(int, 123, 123),
			(int, None, None),
			(float, 123.0, 123),
			(float, 123.45, 123),
			(float, 0.45, 0),
			(float, -0.45, 0),
			(float, -123.45, -123),
			(float, None, None),
			(str, '-123', -123),
			(str, '123', 123),
			(str, '0', 0),
			(str, None, None),
		]
		for src_type, value, expected in cases:
			expr = Cast(ValueExpression(value, src_type, nullable=True), int)
			self.assertEqual(expr.evaluate([]), expected)

	def test_should_return_error_casting_invalid_string_to_integer(self):
		for s in ('asf', '3.14', '0x123', 'unknown', 'null', ''):
			expr = Cast(ValueExpression(s, str, nullable=True), int)
			with self.assertRaisesRegex(TypeError, 'invalid integer'):
				expr.evaluate([])

	def test_should_return_error_casting_bool_to_float(self):
		for b in (False, True, None):
			with self.assertRaisesRegex(TypeError, 'Cannot cast'):
				Cast(ValueExpression(b, bool, nullable=True), float)

	def test_should_cast_to_float(self):
		cases = [
			(int, 123, 123.0),
			(int, None, None),
			(float, 123.0, 123.0),
			(float, 123.45, 123.45),
			(float, None, None),
			(str, '-123', -123),
			(str, '123', 123),
			(str, '0', 0),
			(str, '-123.45', -123.45),
			(str, None, None),
		]
		for src_type, value, expected in cases:
			expr = Cast(ValueExpression(value, src_type, nullable=True), float)
			self.assertEqual(expr.evaluate([]), expected)

	def test_should_return_error_casting_invalid_string_to_float(self):
		for s in ('asf', '0x123', 'unknown', 'null', ''):
			expr = Cast(ValueExpression(s, str, nullable=True), float)
			with self.assertRaisesRegex(TypeError, 'invalid float'):
				expr.evaluate([])

	def test_should_cast_to_str(self):
		cases = [
			(bool, False, 'false'),
			(bool, True, 'true'),
			(bool, None, None),
			(int, 123, '123'),
			(int, -123, '-123'),
			(int, None, None),
			(float, 123.0, '123.0'),
			(float, 123.45, '123.45'),
			(float, None, None),
			(str, '', ''),
			(str, 'abc', 'abc'),
			(str, None, None),
		]
		for src_type, value, expected in cases:
			expr = Cast(ValueExpression(value, src_type, nullable=True), str)
			self.assertEqual(expr.evaluate([]), expected)

class TestAnd(unittest.TestCase):
	def test_should_evaluate_for_booleans(self):
		cases= [
			(None, None, None),
			(None, False, False),
			(None, True, None),
			(False, None, False),
			(False, False, False),
			(False, True, False),
			(True, None, None),
			(True, False, False),
			(True, True, True),
		]
		for lhs, rhs, expected_result in cases:
			expr = And(ValueExpression(lhs, bool),
					ValueExpression(rhs, bool))
			display = '%r AND %r' % (lhs, rhs)
			self.assertEqual(expr.evaluate([]), expected_result, msg=display)

	def test_should_return_type_error_for_non_booleans(self):
		for incorrect_type in (int, float, str):
			with self.assertRaisesRegex(TypeError, 'boolean'):
				And(ValueExpression(None, incorrect_type, nullable=True),
					ValueExpression(True, bool))
			with self.assertRaisesRegex(TypeError, 'boolean'):
				And(ValueExpression(True, bool),
					ValueExpression(None, incorrect_type, nullable=True))

class TestOr(unittest.TestCase):
	def test_should_evaluate_for_booleans(self):
		cases= [
			(None, None, None),
			(None, False, None),
			(None, True, True),
			(False, None, None),
			(False, False, False),
			(False, True, True),
			(True, None, True),
			(True, False, True),
			(True, True, True),
		]
		for lhs, rhs, expected_result in cases:
			expr = Or(ValueExpression(lhs, bool),
					ValueExpression(rhs, bool))
			display = '%r OR %r' % (lhs, rhs)
			self.assertEqual(expr.evaluate([]), expected_result, msg=display)

	def test_should_return_type_error_for_non_booleans(self):
		for incorrect_type in (int, float, str):
			with self.assertRaisesRegex(TypeError, 'boolean'):
				Or(ValueExpression(None, incorrect_type, nullable=True),
					ValueExpression(True, bool))
			with self.assertRaisesRegex(TypeError, 'boolean'):
				Or(ValueExpression(True, bool),
					ValueExpression(None, incorrect_type, nullable=True))

class TestArithmetic(unittest.TestCase):
	def test_should_return_error_for_non_numeric_types(self):
		for incorrect_type in (bool, str):
			with self.assertRaisesRegex(TypeError, 'numeric'):
				Arithmetic('+',
					ValueExpression(None, incorrect_type, nullable=True),
					ValueExpression(0, int))
			with self.assertRaisesRegex(TypeError, 'numeric'):
				Arithmetic('+',
					ValueExpression(0, int),
					ValueExpression(None, incorrect_type, nullable=True))

	def test_should_produce_numeric_type(self):
		cases = [
			('*', 7, 3, 7*3, int),
			('*', 7, 3.1, 7*3.1, float),
			('*', 7.1, 3, 7.1*3, float),
			('*', 7.4, 3.2, 7.4*3.2, float),
			('/', 7, 3, 2, int), # integer division
			('/', 7, 3.1, 7/3.1, float),
			('/', 7.1, 3, 7.1/3, float),
			('/', 7.4, 3.2, 7.4/3.2, float),
			('%', 7, 3, 7%3, int),
			('%', 7, 3.1, 7%3.1, float),
			('%', 7.1, 3, 7.1%3, float),
			('%', 7.4, 3.2, 7.4%3.2, float),
			('+', 7, 3, 7+3, int),
			('+', 7, 3.1, 7+3.1, float),
			('+', 7.1, 3, 7.1+3, float),
			('+', 7.4, 3.2, 7.4+3.2, float),
			('-', 7, 3, 7-3, int),
			('-', 7, 3.1, 7-3.1, float),
			('-', 7.1, 3, 7.1-3, float),
			('-', 7.4, 3.2, 7.4-3.2, float)
		]
		for op, lhs, rhs, expected_value, expected_type in cases:
			expr = Arithmetic(op, Constant(lhs), Constant(rhs))
			display = '%d %s %d' % (lhs, op, rhs)
			self.assertEqual(expr.value_type(), expected_type, msg=display)
			self.assertEqual(expr.evaluate([]), expected_value, msg=display)

	def test_should_do_integer_division(self):
		expr = Arithmetic('/', ValueExpression(7, int), ValueExpression(3, int))
		self.assertEqual(expr.value_type(), int)
		self.assertEqual(expr.evaluate([]), 2)

class TestUnaryMinus(unittest.TestCase):
	def test_should_return_error_for_non_numeric_type(self):
		for incorrect_type in (bool, str):
			with self.assertRaisesRegex(TypeError, 'must be numeric'):
				UnaryMinus(ValueExpression(None, incorrect_type, nullable=True))

	def test_should_return_same_numeric_type_as_operand(self):
		cases = [
			(int, 123, -123),
			(int, -123, 123),
			(float, 3.14, -3.14),
			(int, None, None),
			(float, None, None)
		]
		for operand_type, operand, expected_value in cases:
			expr = UnaryMinus(
					ValueExpression(operand, operand_type, nullable=True))
			self.assertEqual(expr.value_type(), operand_type)
			self.assertEqual(expr.evaluate([]), expected_value)

class TestLogicalNot(unittest.TestCase):
	def test_should_return_error_for_non_boolean_types(self):
		for incorrect_type in (int, float, str):
			with self.assertRaisesRegex(TypeError, 'must be boolean'):
				LogicalNot(ValueExpression(None, incorrect_type, nullable=True))

	def test_should_evaluate_to_expected_values(self):
		cases = [
			(False, True),
			(True, False),
			(None, None)
		]
		for operand, expected_value in cases:
			expr = LogicalNot(
					ValueExpression(operand, bool, nullable=True))
			self.assertEqual(expr.value_type(), bool)
			self.assertEqual(expr.evaluate([]), expected_value)

class TestIsNull(unittest.TestCase):
	def test_should_evaluate_to_expected_values(self):
		cases = [
			(bool, True, False),
			(bool, None, True),
			(int, 123, False),
			(int, None, True),
			(float, 3.14, False),
			(float, None, True),
			(str, '', False),
			(str, 'hello', False),
			(str, None, True),
		]
		for operand_type, operand, expected_value in cases:
			expr = IsNull(
					ValueExpression(operand, operand_type, nullable=True))
			self.assertEqual(expr.value_type(), bool)
			self.assertFalse(expr.nullable())
			self.assertEqual(expr.evaluate([]), expected_value)

class TestIsNotNull(unittest.TestCase):
	def test_should_evaluate_to_expected_values(self):
		cases = [
			(bool, True, True),
			(bool, None, False),
			(int, 123, True),
			(int, None, False),
			(float, 3.14, True),
			(float, None, False),
			(str, '', True),
			(str, 'hello', True),
			(str, None, False),
		]
		for operand_type, operand, expected_value in cases:
			expr = IsNotNull(
					ValueExpression(operand, operand_type, nullable=True))
			self.assertEqual(expr.value_type(), bool)
			self.assertFalse(expr.nullable())
			self.assertEqual(expr.evaluate([]), expected_value)

class TestSelection(unittest.TestCase):
	def test_should_return_error_for_non_boolean_expression(self):
		for incorrect_type in (int, float, str):
			with self.assertRaisesRegex(TypeError, 'Predicate'):
				Selection(MaterialRelation([Column('x', int)]),
				ValueExpression(None, incorrect_type))

	def test_should_have_same_columns_as_input(self):
		relation = Selection(MaterialRelation([
			Column('name', str, nullable=False),
			Column('age', int, nullable=True)
		], name='Users'),
		Constant(False))

		self.assertEqual(len(relation.columns), 2)
		self.assertEqual(relation.columns[0].name, 'name')
		self.assertEqual(relation.columns[0].type, str)
		self.assertEqual(relation.columns[0].nullable, False)
		self.assertEqual(relation.columns[0].index, 0)
		self.assertEqual(relation.columns[1].name, 'age')
		self.assertEqual(relation.columns[1].type, int)
		self.assertEqual(relation.columns[1].nullable, True)
		self.assertEqual(relation.columns[1].index, 1)

		self.assertIsNone(relation.name)
		relation.set_name('NoUsers')
		self.assertEqual(relation.name, 'NoUsers')

	def test_should_return_elements_matching_predicate(self):
		relation = MaterialRelation([
			Column('name', str), Column('age', int), Column('female', bool)
		])
		relation.insert(('Alice', 25, True))
		relation.insert(('Bob', 24, False))
		relation.insert(('Eve', 21, True))
		relation.insert(('Mallory', 35, True))

		selection = Selection(relation,
						Or(Comparison('<', Attribute(relation.columns[1]),
											Constant(24)),
							LogicalNot(Attribute(relation.columns[2]))))
		self.assertEqual(list(selection),
						[('Bob', 24, False), ('Eve', 21, True)])

	def test_should_not_return_elements_where_predicate_is_null(self):
		relation = MaterialRelation([
			Column('name', str), Column('age', int, nullable=True),
		])
		relation.insert(('Alice', 25))
		relation.insert(('Bob', 24))
		relation.insert(('Eve', 21))
		relation.insert(('Mallory', None))

		selection = Selection(relation, Comparison('<',
											Attribute(relation.columns[1]),
											Constant(24)))
		self.assertEqual(list(selection), [('Eve', 21)])

class TestGeneralizedProjection(unittest.TestCase):
	def test_should_have_schema(self):
		relation = MaterialRelation([
			Column('name', str), Column('age', int, nullable=True),
			Column('female', bool, nullable=False)
		])
		projection = GeneralizedProjection(relation, [
			Arithmetic('-', Constant(2021), Attribute(relation.columns[1])),
			LogicalNot(Attribute(relation.columns[2])),
		])
		self.assertEqual(len(projection.columns), 2)
		self.assertEqual(projection.columns[0].type, int)
		self.assertEqual(projection.columns[0].nullable, True)
		self.assertEqual(projection.columns[0].index, 0)
		self.assertEqual(projection.columns[1].type, bool)
		self.assertEqual(projection.columns[1].nullable, False)
		self.assertEqual(projection.columns[1].index, 1)

	def test_should_have_expected_column_names(self):
		relation = MaterialRelation([
			Column('name', str), Column('age', int)
		])
		projection = GeneralizedProjection(relation, [
			Attribute(relation.columns[0]),
			Cast(Attribute(relation.columns[1]), float),
			Arithmetic('-', Constant(2021), Attribute(relation.columns[1])),
		])
		self.assertEqual(len(projection.columns), 3)
		self.assertEqual(projection.columns[0].name, 'name')
		self.assertEqual(projection.columns[1].name, 'age')
		self.assertEqual(projection.columns[2].name, None)

	def test_should_generate_output(self):
		relation = MaterialRelation([
			Column('name', str), Column('age', int, nullable=True),
			Column('female', bool, nullable=False)
		])
		relation.insert(('Alice', 25, True))
		relation.insert(('Bob', 24, False))
		relation.insert(('Eve', 21, True))
		relation.insert(('Mallory', 35, True))

		projection = GeneralizedProjection(relation, [
			Arithmetic('-', Constant(2021), Attribute(relation.columns[1])),
			LogicalNot(Attribute(relation.columns[2]))
		])

		self.assertEqual(list(projection),
					[(1996, False), (1997, True), (2000, False), (1986, False)])

class TestNextValue(unittest.TestCase):
	def test(self):
		a = [1,2,3]
		it = a.__iter__()

		self.assertEqual(next_value(it), 1)
		self.assertEqual(next_value(it), 2)
		self.assertEqual(next_value(it), 3)
		self.assertIsNone(next_value(it))
		self.assertIsNone(next_value(it))

class TestRemoveDuplicates(unittest.TestCase):
	def test_empty_stream(self):
		self.assertEqual(list(remove_duplicates([].__iter__())), [])

	def test_removes_duplicates(self):
		cases = [
			([], []),
			([1, 2, 3], [1, 2, 3]),
			([1, 1, 2, 2, 3, 3], [1, 2, 3]),
		]
		for dupes, deduped in cases:
			self.assertEqual(list(remove_duplicates(dupes.__iter__())), deduped)

class TestStreamUnion(unittest.TestCase):
	def test(self):
		cases = [
			([], [1, 2, 3, 3], [1, 2, 3, 3]),
			([1, 3], [2, 4], [1, 2, 3, 4]),
			([1, 3], [3, 4], [1, 3, 3, 4]),
		]
		for a, b, u in cases:
			self.assertEqual(list(stream_union(a.__iter__(), b.__iter__())), u)
			self.assertEqual(list(stream_union(b.__iter__(), a.__iter__())), u)

class TestStreamIntersection(unittest.TestCase):
	def test(self):
		cases = [
			([], [1, 2, 3, 3], []),
			([1, 3], [2, 4], []),
			([1, 3], [3, 4], [3]),
			([1, 3, 3, 4], [3, 4], [3, 3, 3, 4]),
		]
		for a, b, i in cases:
			self.assertEqual(
				list(stream_intersection(a.__iter__(), b.__iter__())), i)
			self.assertEqual(
				list(stream_intersection(b.__iter__(), a.__iter__())), i)

class TestStreamDifference(unittest.TestCase):
	def test(self):
		cases = [
			([], [1, 2, 3, 3], []),
			([1, 2, 3, 3], [], [1, 2, 3, 3]),
			([1, 3], [2, 4], [1, 3]),
			([1, 1, 3, 3], [3, 4], [1]),
		]
		for a, b, d in cases:
			self.assertEqual(
				list(stream_difference(a.__iter__(), b.__iter__())), d)

class TestStreamIntersection(unittest.TestCase):
	def test(self):
		lhs = [1, 1, 1, 3, 4, 5, 5, 6].__iter__()
		rhs = [1, 4, 6, 7].__iter__()
		self.assertEqual(list(stream_intersection(lhs, rhs)),
				[1, 1, 1, 1, 4, 4, 6, 6])

class TestStreamDifference(unittest.TestCase):
	def test(self):
		lhs = [1, 1, 1, 3, 4, 5, 5, 6].__iter__()
		rhs = [1, 4, 6, 7].__iter__()
		self.assertEqual(list(stream_difference(lhs, rhs)), [3, 5, 5])

class TestCompareTuples(unittest.TestCase):
	def test(self):
		test_cases = [
			((2, 3), (2, 4), True, -1),
			((2, 3), (2, 3), True, 0),
			((2, 4), (2, 3), True, 1),
			((2, 3), (2, None), True, -1),
			((2, 3), (2, None), False, 1),
			((2, None), (2, 3), True, 1),
			((2, None), (2, 3), False, -1),
			((None, 2), (None, 2), True, 0),
			((None, 2), (None, 3), True, -1),
			((None, 3), (None, 2), True, 1),
		]
		for lhs, rhs, nulls_last, expected in test_cases:
			self.assertEqual(compare_tuples(lhs, rhs, nulls_last), expected,
				msg='%r < %r %r' % (lhs, rhs, nulls_last))

class TestSort(unittest.TestCase):
	def test_should_sort_ascending_by_default(self):
		relation = MaterialRelation([
			Column('name', str), Column('age', int),
		])
		relation.insert(('Mallory', 35))
		relation.insert(('Alice', 13))
		relation.insert(('Bob', 24))
		relation.insert(('Alice', 25))

		ordered = Sort(relation)
		self.assertEqual(list(ordered),
			[('Alice', 13), ('Alice', 25), ('Bob', 24), ('Mallory', 35)])

	def test_descending(self):
		relation = MaterialRelation([Column('name', str)])
		relation.insert(('Mallory',))
		relation.insert(('Alice',))
		relation.insert(('Bob',))

		ordered = Sort(relation, descending=True)
		self.assertEqual(list(ordered), [('Mallory',), ('Bob',), ('Alice',)])

	def test_should_sort_by_attributes(self):
		relation = MaterialRelation([
			Column('name', str), Column('age', int),
		])
		relation.insert(('Mallory', 35))
		relation.insert(('Alice', 13))
		relation.insert(('Bob', 24))
		relation.insert(('Alice', 25))

		ordered = Sort(relation,
						sort_key=[relation.columns[1], relation.columns[0]])
		self.assertEqual(list(ordered),
			[('Alice', 13), ('Bob', 24), ('Alice', 25), ('Mallory', 35)])

	def test_sort_orders_null_first(self):
		relation = MaterialRelation([
			Column('age', int, nullable=True),
		])
		relation.insert((35,))
		relation.insert((13,))
		relation.insert((None,))
		relation.insert((25,))

		ordered = Sort(relation, nulls_last=False)
		self.assertEqual(list(ordered), [(None,), (13,), (25,), (35,)])

	def test_sort_orders_null_last_by_default(self):
		relation = MaterialRelation([
			Column('age', int, nullable=True),
		])
		relation.insert((35,))
		relation.insert((13,))
		relation.insert((None,))
		relation.insert((25,))

		ordered = Sort(relation)
		self.assertEqual(list(ordered), [(13,), (25,), (35,), (None,)])

	def test_should_have_same_columns_as_input(self):
		relation = Sort(MaterialRelation([
			Column('name', str, nullable=False),
			Column('age', int, nullable=True)
		], name='Users'))

		self.assertEqual(len(relation.columns), 2)
		self.assertEqual(relation.columns[0].name, 'name')
		self.assertEqual(relation.columns[0].type, str)
		self.assertEqual(relation.columns[0].nullable, False)
		self.assertEqual(relation.columns[0].index, 0)
		self.assertEqual(relation.columns[1].name, 'age')
		self.assertEqual(relation.columns[1].type, int)
		self.assertEqual(relation.columns[1].nullable, True)
		self.assertEqual(relation.columns[1].index, 1)

		self.assertIsNone(relation.name)

class TestUnion(unittest.TestCase):
	def test_should_return_error_for_varying_tuple_length(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs = MaterialRelation([Column('a', str)])
		with self.assertRaisesRegex(ValueError, 'number of columns'):
			Union(lhs, rhs)

	def test_should_return_error_if_relations_have_different_column_types(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs = MaterialRelation([Column('a', str), Column('b', str)])
		with self.assertRaisesRegex(ValueError, 'column types'):
			Union(lhs, rhs)

	def test_should_return_duplicate_tuples_for_union_all(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		lhs.insert(('au', 123))
		lhs.insert(('ca', 456))
		lhs.insert(('ca', 456))

		rhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs.insert(('fr', 123))
		rhs.insert(('ca', 456))
		rhs.insert(('ch', 789))

		union = Union(lhs, rhs, distinct=False)
		self.assertEqual(list(union), [('au', 123), ('ca', 456), ('ca', 456),
			('ca', 456), ('ch', 789), ('fr', 123)])

	def test_should_omit_duplicate_tuples_for_union_distinct(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		lhs.insert(('au', 123))
		lhs.insert(('ca', 456))
		lhs.insert(('ca', 456))

		rhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs.insert(('fr', 123))
		rhs.insert(('ca', 456))
		rhs.insert(('ch', 789))

		union = Union(lhs, rhs, distinct=True)
		self.assertEqual(list(union),
			[('au', 123), ('ca', 456), ('ch', 789), ('fr', 123)])

	def test_should_merge_relations_with_different_column_names(self):
		lhs = MaterialRelation([Column('a', str)])
		lhs.insert(('au',))

		rhs = MaterialRelation([Column('b', str)])
		rhs.insert(('fr',))
		union = Union(lhs, rhs)
		self.assertEqual(list(union), [('au',), ('fr',)])

	def test_should_use_first_relations_column_names_for_output(self):
		lhs = MaterialRelation([Column('a', str), Column('x', int)])
		rhs = MaterialRelation([Column('b', str), Column('y', int)])
		union = Union(lhs, rhs)
		self.assertEqual(len(union.columns), 2)
		self.assertEqual(union.columns[0].name, 'a')
		self.assertEqual(union.columns[0].index, 0)
		self.assertEqual(union.columns[1].name, 'x')
		self.assertEqual(union.columns[1].index, 1)

	def test_should_be_nullable_when_any_input_relation_is_nullable(self):
		lhs = MaterialRelation([
			Column('a', str, nullable=False),
			Column('b', int, nullable=False),
		])
		rhs = MaterialRelation([
			Column('a', str, nullable=False),
			Column('b', int, nullable=True),
		])
		union = Union(lhs, rhs)
		self.assertFalse(union.columns[0].nullable)
		self.assertTrue(union.columns[1].nullable)

class TestIntersection(unittest.TestCase):
	def test_should_return_duplicate_tuples_for_intersection_all(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		lhs.insert(('au', 123))
		lhs.insert(('ca', 456))
		lhs.insert(('ca', 456))

		rhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs.insert(('fr', 123))
		rhs.insert(('ca', 456))
		rhs.insert(('ch', 789))

		intersection = Intersection(lhs, rhs, distinct=False)
		self.assertEqual(list(intersection), [('ca', 456), ('ca', 456),
			('ca', 456)])

	def test_should_omit_duplicate_tuples_for_intersection_distinct(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		lhs.insert(('au', 123))
		lhs.insert(('ca', 456))
		lhs.insert(('ca', 456))
		lhs.insert(('ch', 789))

		rhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs.insert(('fr', 123))
		rhs.insert(('ca', 456))
		rhs.insert(('ch', 789))

		intersection = Intersection(lhs, rhs, distinct=True)
		self.assertEqual(list(intersection), [('ca', 456), ('ch', 789)])
		self.assertTrue(intersection.columns[1].nullable)

class TestDifference(unittest.TestCase):
	def test_should_return_duplicate_tuples_for_difference_all(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		lhs.insert(('au', 123))
		lhs.insert(('au', 123))
		lhs.insert(('ca', 456))

		rhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs.insert(('fr', 123))
		rhs.insert(('ca', 456))
		rhs.insert(('ch', 789))

		difference = Difference(lhs, rhs, distinct=False)
		self.assertEqual(list(difference), [('au', 123), ('au', 123)])

	def test_should_omit_duplicate_tuples_for_difference_distinct(self):
		lhs = MaterialRelation([Column('a', str), Column('b', int)])
		lhs.insert(('au', 123))
		lhs.insert(('au', 123))
		lhs.insert(('ca', 456))

		rhs = MaterialRelation([Column('a', str), Column('b', int)])
		rhs.insert(('fr', 123))
		rhs.insert(('ca', 456))
		rhs.insert(('ch', 789))

		difference = Difference(lhs, rhs, distinct=True)
		self.assertEqual(list(difference), [('au', 123)])

# TODO:
# COUNT(*) - number of input rows
# COUNT(expr) - number of rows expression is not null

class TestGroupBy(unittest.TestCase):
	def test_should_have_correct_schema(self):
		relation = MaterialRelation([Column('a', str, nullable=True),
				Column('b', bool, nullable=False),
				Column('c', int)])
		output = GroupBy(relation, relation.columns[:2], [CountFactory()])

		self.assertEqual(len(output.columns), 3)
		self.assertEqual(output.columns[0].name, 'a')
		self.assertEqual(output.columns[0].type, str)
		self.assertEqual(output.columns[0].nullable, True)
		self.assertEqual(output.columns[0].index, 0)
		self.assertEqual(output.columns[1].name, 'b')
		self.assertEqual(output.columns[1].type, bool)
		self.assertEqual(output.columns[1].nullable, False)
		self.assertEqual(output.columns[1].index, 1)
		self.assertEqual(output.columns[2].name, None)
		self.assertEqual(output.columns[2].type, int)
		self.assertEqual(output.columns[2].nullable, False)
		self.assertEqual(output.columns[2].index, 2)

	# TODO:
	# Function, type, nullable?
	# count, int, false
	# min, * , true
	# max, * , true
	# sum, <numeric>, false (default is 0)
	# avg - no tuples, no non-null values, non-numeric expression type

	# sum - null values, non-numeric expression type
	# sum - only valid for numeric types
	#	- what if expression evaluates to null for non-null input?
	#	- SUM(NULL)
	#	- default value of sum is 0 (all null input)
	# min, max - int, str, float, ...
	# count if

	# group by - nth column

# Expressions involving aggregates
# SELECT MAX(x) - MIN(x) FROM R GROUP BY y;

class TestCrossJoin(unittest.TestCase):
	def test(self):
		lhs = MaterialRelation([
			Column('id', int, nullable=False),
			Column('name', str, nullable=False)
		])
		lhs.insert((1, 'Alice'))
		lhs.insert((2, 'Bob'))
		lhs.insert((3, 'Eve'))

		rhs = MaterialRelation([
			Column('x', str, nullable=True),
			Column('y', bool, nullable=True),
			Column('z', float, nullable=False),
		])
		rhs.insert(('x', False, 3.14))
		rhs.insert(('X', True, 2.71))

		product = CrossJoin(lhs, rhs)
		self.assertEqual(len(product.columns), 5)
		self.assertEqual(product.columns[0].name, 'id')
		self.assertEqual(product.columns[0].type, int)
		self.assertFalse(product.columns[0].nullable)
		self.assertEqual(product.columns[0].index, 0)
		self.assertEqual(product.columns[1].name, 'name')
		self.assertEqual(product.columns[1].type, str)
		self.assertFalse(product.columns[1].nullable)
		self.assertEqual(product.columns[1].index, 1)
		self.assertEqual(product.columns[2].name, 'x')
		self.assertEqual(product.columns[2].type, str)
		self.assertTrue(product.columns[2].nullable)
		self.assertEqual(product.columns[2].index, 2)
		self.assertEqual(product.columns[3].name, 'y')
		self.assertEqual(product.columns[3].type, bool)
		self.assertTrue(product.columns[3].nullable)
		self.assertEqual(product.columns[3].index, 3)
		self.assertEqual(product.columns[4].name, 'z')
		self.assertEqual(product.columns[4].type, float)
		self.assertFalse(product.columns[4].nullable)

		self.assertEqual(list(product), [
			(1, 'Alice', 'x', False, 3.14),
			(1, 'Alice', 'X', True, 2.71),
			(2, 'Bob', 'x', False, 3.14),
			(2, 'Bob', 'X', True, 2.71),
			(3, 'Eve', 'x', False, 3.14),
			(3, 'Eve', 'X', True, 2.71),
		])

	def test_should_allow_multiple_columns_with_same_name(self):
		lhs = MaterialRelation([
			Column('a', int),
			Column('b', str)
		])
		lhs.insert((1, 'Alice'))
		lhs.insert((2, 'Bob'))
		lhs.insert((3, 'Eve'))

		rhs = MaterialRelation([
			Column('b', str, nullable=True),
			Column('c', bool, nullable=True),
		])
		rhs.insert(('x', False))
		rhs.insert(('X', True))

		product = CrossJoin(lhs, rhs)

		self.assertEqual(len(product.columns), 4)
		self.assertEqual(product.columns[0].name, 'a')
		self.assertEqual(product.columns[1].name, 'b')
		self.assertEqual(product.columns[2].name, 'b')
		self.assertEqual(product.columns[3].name, 'c')

		self.assertEqual(list(product), [
			(1, 'Alice', 'x', False),
			(1, 'Alice', 'X', True),
			(2, 'Bob', 'x', False),
			(2, 'Bob', 'X', True),
			(3, 'Eve', 'x', False),
			(3, 'Eve', 'X', True),
		])

# TODO:
# - expression in select predicate, generalized projection, or aggregation
#   references columns not in the input relation
# nulls
# implicit conversions
# TODO: test type and nullability attributes

# TODO:
# attributes - value type, nullability
# different types - throw error for now
# implicit cast bool to int to float
# comparison - bool -> int -> float -> string
# nullability - short circuiting

if __name__ == '__main__':
	unittest.main()


from threading import Thread

MSGLEN = 500

class Sender(Thread):
    def __init__(self, queue, socket):
        super().__init__()
        self.queue = queue
        self.socket = socket

    def send_to_server(self, message):
        total_sent = 0
        message = message.zfill(MSGLEN)
        while total_sent < MSGLEN:
            sent = self.socket.send(message[total_sent:].encode())
            if sent == 0:
                raise RuntimeError("Socket connection broken")
            total_sent = total_sent + sent

    def run(self):
        while True:
            message = self.queue.get()
            if message == "end client":
                break
            self.send_to_server(message)
        self.socket.close()


"""Reviewers and ratings: Number of Submissions/Revisions over time"""
from asyncio import gather, run
from typing import Tuple, Type, Union

from matplotlib import pyplot as plt
from matplotlib.axes import Axes
from pandas import DataFrame
from seaborn import barplot, color_palette
from sqlalchemy import func
from sqlalchemy.future import select
from sqlalchemy.sql.selectable import Select
from uvloop import install

from cmt_statistics_tool.statistics import format_sort_track, get_data, plot_df
from cmt_statistics_tool.tables import Revision, Submission

def statement(paper: Union[Type[Submission], Type[Revision]]) -> Select:
    return (
        select(paper.track_name, func.count())
        .group_by(paper.track_name)
        .order_by(paper.track_name)
    )

async def submission() -> DataFrame:
    df = DataFrame(await get_data(statement(Submission))).rename(
        columns={0: "Track", 1: "Count"}
    )
    return format_sort_track(df, "Track")

def plot_submission(df: DataFrame, ax: Axes) -> None:
    barplot(x="Track", y="Count", data=df, color=color_palette()[0], ax=ax)
    ax.set_title("Number of submissions over time")

async def revision() -> DataFrame:
    df = DataFrame(await get_data(statement(Revision))).rename(
        columns={0: "Track", 1: "Count"}
    )
    return format_sort_track(df, "Track", True)

def plot_revision(df: DataFrame, ax: Axes) -> None:
    barplot(x="Track", y="Count", data=df, color=color_palette()[0], ax=ax)
    ax.set_title("Number of revisions over time")

def plot_both(df: DataFrame, ax: Axes) -> None:
    barplot(x="Track", y="Count", hue="Type", data=df, ax=ax)
    ax.set_title("Number of papers over time")

async def main() -> Tuple[DataFrame, DataFrame]:
    return await gather(submission(), revision())

if __name__ == "__main__":
    install()
    s_df, r_df = run(main())
    plot_df(s_df, plot_submission)
    plt.savefig("plots/01_03_submission.png")
    plot_df(r_df, plot_revision)
    plt.savefig("plots/01_03_revision.png")
    print(s_df, r_df, sep="\n")
    plot_df(
        s_df.set_index("Track")
        .rename(columns={"Count": "Original Submission"})
        .join(
            r_df.set_index("Track").rename(columns={"Count": "Revision"}), how="outer"
        )
        .fillna(0)
        .astype(int)
        .reset_index()
        .melt(id_vars=["Track"], var_name="Type", value_name="Count"),
        plot_both,
    )
    plt.savefig("plots/01_03_both.png")


from typing import List, Any
from unittest.mock import patch

import torch
import pytest

from torchsim.core import get_float, SMALL_CONSTANT
from torchsim.core.memory.tensor_creator import AllocatingCreator
from torchsim.core.models.expert_params import ExpertParams, DEFAULT_CONTEXT_PRIOR, NUMBER_OF_CONTEXT_TYPES
from torchsim.core.models.temporal_pooler import TPFlock, UntrainedForwardProcessFactory, \
    TrainedForwardProcessFactory, TPOutputProjection
from torchsim.core.utils.tensor_utils import same
from tests.testing_utils import copy_or_fill_tensor

def create_tp_flock(n_cluster_centers=3, flock_size=1, seq_length=3, buffer_size=10, batch_size=5,
                    seq_lookahead=1, n_frequent_seq=4, max_encountered_seq=20,
                    forgetting_limit=40, learning_period=1, incoming_context_size=1,
                    context_prior=DEFAULT_CONTEXT_PRIOR, exploration_probability=0, follow_goals=False,
                    n_providers=1,
                    trained_forward_factory=None, untrained_forward_factory=None, enable_learning=True, device='cuda'):
    params = ExpertParams()
    params.flock_size = flock_size
    params.n_cluster_centers = n_cluster_centers
    params.flock_size = flock_size

    tp_params = params.temporal
    tp_params.incoming_context_size = incoming_context_size
    tp_params.seq_length = seq_length
    tp_params.buffer_size = buffer_size
    tp_params.batch_size = batch_size
    tp_params.seq_lookahead = seq_lookahead
    tp_params.n_frequent_seqs = n_frequent_seq
    tp_params.max_encountered_seqs = max_encountered_seq
    tp_params.forgetting_limit = forgetting_limit
    tp_params.learning_period = learning_period
    tp_params.context_prior = context_prior
    tp_params.exploration_probability = exploration_probability
    tp_params.follow_goals = follow_goals
    tp_params.enable_learning = enable_learning
    tp_params.n_providers = n_providers
    tp_params.frustration_threshold = 10000
    tp_params.follow_goals = False
    tp_params.compute_backward_pass = True

    return TPFlock(params, creator=AllocatingCreator(device),
                   trained_forward_factory=trained_forward_factory,
                   untrained_forward_factory=untrained_forward_factory)

# TODO (Test): add a test for follow_goals=True
@pytest.mark.slow
def test_forward_learn_no_context():
    device = 'cuda'
    float_dtype = get_float(device)
    flock = create_tp_flock(flock_size=2, seq_length=3, buffer_size=1000, batch_size=20, max_encountered_seq=800,
                            incoming_context_size=4, exploration_probability=0)

    # experts:
    #   ex1 receives the same input all the time - doesn't learn
    #   ex2 receives a cycling sequence - this one should learn appropriately

    seqs = torch.tensor([[[1, 0, 0], [1, 0, 0]],
                         [[1, 0, 0], [0, 1, 0]],
                         [[1, 0, 0], [0, 0, 1]],
                         [[1, 0, 0], [0, 1, 0]]], dtype=float_dtype, device=device)

    context = torch.tensor([[[[1, 0, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0]]], [[[1, 0, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0]]]], dtype=float_dtype,
                           device=device)

    # If the forward pass doesnt run, the outputs should always be zero for ex1
    expected_actions = torch.tensor([[[0., 0., 0.], [0, 1, 0]],
                                     [[0., 0., 0.], [0, 0, 1]],
                                     [[0., 0., 0.], [0, 1, 0]],
                                     [[0., 0., 0.], [1, 0, 0]]], dtype=float_dtype, device=device)

    expected_projections = torch.tensor([[[0.3333, 0.3333, 0.3333], [0.5, 0.5, 0]],
                                         [[0.3333, 0.3333, 0.3333], [0.25, 0.5, 0.25]],
                                         [[0.3333, 0.3333, 0.3333], [0, 0.5, 0.5]],
                                         [[0.3333, 0.3333, 0.3333], [0.25, 0.5, 0.25]]], dtype=float_dtype,
                                        device=device)

    expected_buffer_total_data_written = torch.tensor([1, 20], dtype=torch.int64,
                                                      device=device)
    # TODO (Test): check other intermediate tensors and values, buffers, etc.

    iterations = 20
    for k in range(iterations):
        cluster_data = seqs[k % 4]
        context_data = context
        flock.forward_learn(cluster_data, context_data, input_rewards=None)

    assert same(expected_buffer_total_data_written, flock.buffer.total_data_written)

    # Run through the sequences yet again and no grab the outputs each time for flock 2
    for k in range(4):
        flock.forward_learn(seqs[k], context, input_rewards=None)
        # Check that the predicted outputs are what we expect
        assert same(expected_actions[k], flock.action_outputs, eps=1e-2)
        # Check that the output projection is what we expect
        assert same(expected_projections[k], flock.projection_outputs, eps=1e-3)

def test_forward_learn_context():
    """Tests if context helps to disambiguate situations."""
    device = 'cuda'
    float_dtype = get_float(device)
    batch_size = 5
    flock = create_tp_flock(flock_size=2, seq_length=2, buffer_size=1000, batch_size=batch_size,
                            max_encountered_seq=800, n_frequent_seq=50,
                            incoming_context_size=2, learning_period=4, context_prior=SMALL_CONSTANT,
                            forgetting_limit=1)

    # experts:
    #   ex1 does not receive the context during learning, so can't learn to predict correctly
    #   ex2 receives a cycling sequence and useful context - this one should learn appropriately

    seqs = torch.tensor([[[1, 0, 0], [1, 0, 0]],
                         [[0, 1, 0], [0, 1, 0]],
                         [[0, 0, 1], [0, 0, 1]],
                         [[0, 1, 0], [0, 1, 0]]], dtype=float_dtype, device=device)

    contexts = torch.tensor([[[[[1, 0], [0, 0], [0, 0]]], [[[1, 0], [0, 0], [0, 0]]]],
                             [[[[1, 0], [0, 0], [0, 0]]], [[[1, 0], [0, 0], [0, 0]]]],
                             [[[[1, 0], [0, 0], [0, 0]]], [[[1, 0], [0, 0], [0, 0]]]],
                             [[[[1, 0], [0, 0], [0, 0]]], [[[0, 1], [0, 0], [0, 0]]]]], dtype=float_dtype, device=device)

    # In ex1, the context doesnt help, in ex2, the context does
    expected_actions = torch.tensor([[[0, 1, 0], [0, 1, 0]],
                                     [[0.5, 0, 0.5], [0, 0, 1]],
                                     [[0, 1, 0], [0, 1, 0]],
                                     [[0.5, 0, 0.5], [1, 0, 0]]], dtype=float_dtype, device=device)

    expected_projections = torch.tensor([[[0.6666, 0.3333, 0.], [0.6666, 0.3333, 0.]],
                                         [[0.1666, 0.6666, 0.1666], [0, 0.6666, 0.3333]],
                                         [[0., 0.3333, 0.6666], [0., 0.3333, 0.6666]],
                                         [[0.1666, 0.6666, 0.1666], [0.3333, 0.6666, 0.]]], dtype=float_dtype,
                                        device=device)

    # need to learn at least 2 times because contexts are not extracted when sequence is new. And they are regarded
    #  as uniform context so we need even more to get to a sharp probability
    iterations = batch_size * 4
    for k in range(iterations):
        cluster_data = seqs[k % 4]
        context_data = contexts[k % 4]
        flock.forward_learn(cluster_data, context_data, input_rewards=None)

    # Run through the sequences yet again and no grab the outputs each time for flock 2
    for k in range(4):
        flock.forward_learn(seqs[k], contexts[k], input_rewards=None)
        # Check that the predicted outputs are what we expect
        assert same(expected_actions[k], flock.action_outputs, eps=1e-2)
        # Check that the output projection is what we expect
        assert same(expected_projections[k], flock.projection_outputs, eps=3e-2)

@pytest.mark.skip("Test obseleted")
def test_forward_learn_reward():
    """Tests if the agent tries to reach the reward.

    There are two sequences: 0, 1, 2 and 0, 3, 4.
    There can be a reward associated with cluster 3 or punishment associated with cluster 2.
    """
    device = 'cuda'
    float_dtype = get_float(device)
    data_length = 6
    batch_size = data_length + 2
    flock_size = 3
    n_cluster_centers = 5
    flock = create_tp_flock(flock_size=flock_size, seq_length=3, seq_lookahead=2, buffer_size=100,
                            batch_size=batch_size, max_encountered_seq=20, n_frequent_seq=10,
                            n_cluster_centers=n_cluster_centers, incoming_context_size=1,
                            learning_period=data_length, context_prior=0.001, forgetting_limit=1,
                            follow_goals=True)

    seqs = torch.tensor([[[1, 0, 0, 0, 0]],
                         [[0, 1, 0, 0, 0]],
                         [[0, 0, 1, 0, 0]],
                         [[1, 0, 0, 0, 0]],
                         [[0, 0, 0, 1, 0]],
                         [[0, 0, 0, 0, 1]]],
                        dtype=float_dtype, device=device).expand(data_length, flock_size, n_cluster_centers)

    # context does not help in this case
    contexts = torch.tensor([[[[1], [0.5]]]],
                            dtype=float_dtype, device=device).expand(data_length, flock_size, NUMBER_OF_CONTEXT_TYPES,
                                                                     1)

    # experts:
    #   ex1 does not receive the reward during learning, so can't learn to reach it
    #   ex2 receives the positive reward signal - this one should learn to approach it
    #   ex3 receives the negative reward signal (punishment) - this one should learn to avoid it
    rewards = torch.tensor([[[0, 0], [0, 0], [0, 0]],
                            [[0, 0], [0, 0], [0, 0]],
                            [[0, 0], [0, 0], [0, 1]],
                            [[0, 0], [0, 0], [0, 0]],
                            [[0, 0], [0, 0], [0, 0]],
                            [[0, 0], [1, 0], [0, 0]]])

    expected_actions = torch.tensor([[[0, 0.5, 0, 0.5, 0], [0, 0, 0, 1, 0], [0, 0, 0, 1, 0]],
                                     [[0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0]],
                                     [[1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0]],
                                     [[0, 0.5, 0, 0.5, 0], [0, 0, 0, 1, 0], [0, 0, 0, 1, 0]],
                                     [[0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1]],
                                     [[1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0]]], dtype=float_dtype,
                                    device=device)

    # need to learn at least 2 times because contexts are not extracted when sequence is new. And they are regarded
    #  as uniform context so we need even more to get to a sharp probability
    iterations = data_length * 4
    for k in range(iterations):
        cluster_data = seqs[k % data_length]
        context_data = contexts[k % data_length]
        rewards_data = rewards[k % data_length]
        flock.forward_learn(cluster_data, context_data, rewards_data)

    # Run through the sequences yet again and no grab the outputs each time for flock 2
    for k in range(data_length):
        flock.forward_learn(seqs[k], contexts[k], input_rewards=rewards[k])
        # Check that the predicted outputs are what we expect
        assert same(expected_actions[k], flock.action_rewards, eps=2e-2)

@pytest.mark.skip(reason="Ignoring non-default stream tests for now.")
def test_forward_learn_non_default_stream():
    device = 'cuda'
    float_dtype = get_float(device)
    flock = create_tp_flock(flock_size=2, seq_length=3, buffer_size=1000, batch_size=20, max_encountered_seq=800)

    # experts:
    #   ex1 receives the same input all the time - doesn't learn
    #   ex2 receives a cycling sequence - this one should learn appropriately

    seqs = torch.tensor([[[1, 0, 0], [1, 0, 0]],
                         [[1, 0, 0], [0, 1, 0]],
                         [[1, 0, 0], [0, 0, 1]],
                         [[1, 0, 0], [0, 1, 0]]], dtype=float_dtype, device=device)

    # If the forward pass doesnt run, the outputs should always be zero for ex1
    expected_outputs = torch.tensor([[[[0.3333, 0.3333, 0.3333]], [[0, 1, 0]]],
                                     [[[0.3333, 0.3333, 0.3333]], [[0, 0, 1]]],
                                     [[[0.3333, 0.3333, 0.3333]], [[0, 1, 0]]],
                                     [[[0.3333, 0.3333, 0.3333]], [[1, 0, 0]]]], dtype=float_dtype, device=device)

    expected_projections = torch.tensor([[[0.3333, 0.3333, 0.3333], [0.5, 0.5, 0]],
                                         [[0.3333, 0.3333, 0.3333], [0.25, 0.5, 0.25]],
                                         [[0.3333, 0.3333, 0.3333], [0, 0.5, 0.5]],
                                         [[0.3333, 0.3333, 0.3333], [0.25, 0.5, 0.25]]], dtype=float_dtype,
                                        device=device)

    with torch.cuda.stream(torch.cuda.Stream()):
        iterations = 20
        for k in range(iterations):
            data = seqs[k % 4]
            flock.forward_learn(data)

        # Run through the seqs yet again and no grab the outputs each time for flock 2
        for k in range(4):
            flock.forward_learn(seqs[k])
            # Check that the predicted outputs are what we expect
            assert same(expected_outputs[k], flock.action_rewards, eps=1e-4)
            # Check that the output projection is what we expect
            assert same(expected_projections[k], flock.projection_outputs, eps=1e-4)

@pytest.mark.parametrize('enable_learning', [True, False])
def test_forward_learn_enable_learning(enable_learning):
    device = 'cuda'
    float_dtype = get_float(device)
    flock = create_tp_flock(flock_size=1, seq_length=3, buffer_size=1000, batch_size=4, max_encountered_seq=800,
                            incoming_context_size=1, exploration_probability=0, enable_learning=enable_learning)

    seqs = torch.tensor([[[1, 0, 0]],
                         [[0, 1, 0]],
                         [[0, 0, 1]],
                         [[0, 1, 0]]], dtype=float_dtype, device=device)

    initial_seq_occurrences = flock.all_encountered_seq_occurrences.clone()

    iterations = 4
    for k in range(iterations):
        cluster_data = seqs[k % 4]
        flock.forward_learn(cluster_data, input_context=None, input_rewards=None)

    # should be different if enable_learning == True
    assert (not same(initial_seq_occurrences, flock.all_encountered_seq_occurrences)) == enable_learning

# TODO: merge this test with the SP one - actually, the whole determination process could be probablz moved to the
# TODO: base class
@pytest.mark.parametrize("mask, total_data_written, data_since_last_sample, expected_learn_tensor",
                         [(torch.ones(7, dtype=torch.uint8), 0, 0, torch.zeros(7, dtype=torch.uint8)),
                          (torch.ones(7, dtype=torch.uint8), 5, 5, torch.zeros(7, dtype=torch.uint8)),
                          (torch.ones(7, dtype=torch.uint8), 11, 1, torch.zeros(7, dtype=torch.uint8)),
                          (torch.ones(7, dtype=torch.uint8), 15, 5, torch.ones(7, dtype=torch.uint8)),
                          (torch.zeros(7, dtype=torch.uint8), 15, 5, torch.zeros(7, dtype=torch.uint8)),
                          (torch.tensor([0, 1, 1, 0, 0, 0, 1], dtype=torch.uint8), 15, 5,
                           torch.tensor([0, 1, 1, 0, 0, 0, 1], dtype=torch.uint8)),
                          (torch.tensor([0, 1, 1, 1, 0, 1, 1], dtype=torch.uint8),
                           torch.tensor([11, 9, 11, 11, 15, 8, 100], dtype=torch.uint8),
                           torch.tensor([5, 6, 4, 8, 5, 1, 5], dtype=torch.uint8),
                           torch.tensor([0, 0, 0, 1, 0, 0, 1], dtype=torch.uint8))])
def test_determine_learning(mask, total_data_written, data_since_last_sample, expected_learn_tensor):
    """See _test_determine_learning in tests.core.models.spatial_pooler.test_sp_processes."""
    params = ExpertParams()
    params.flock_size = 7
    params.n_cluster_centers = 5

    t_pooler = params.temporal
    t_pooler.batch_size = 10
    t_pooler.learning_period = 5

    flock = TPFlock(params, creator=AllocatingCreator('cpu'))
    copy_or_fill_tensor(flock.buffer.total_data_written, total_data_written)
    copy_or_fill_tensor(flock.buffer.data_since_last_sample, data_since_last_sample)
    learn_tensor = flock._determine_learning(mask)
    assert same(expected_learn_tensor, learn_tensor)

@pytest.mark.slow
def test_trained_untrained():
    flock_size = 2
    n_cluster_centers = 3
    batch_size = 2
    buffer_size = 5
    seq_length = 2
    device = 'cuda'
    float_dtype = get_float(device)

    trained_factory = TrainedForwardProcessFactory()
    untrained_factory = UntrainedForwardProcessFactory()

    # sequence of length 4 of input clusters (3) for each expert in the flock (2)
    input_clusters_seq = torch.tensor([
        [[1, 0, 0], [1, 0, 0]],
        [[0, 1, 0], [1, 0, 0]],
        [[0, 0, 1], [0, 1, 0]],
        [[0, 1, 0], [0, 0, 1]]
    ], dtype=float_dtype, device=device)

    # The first expert:
    # Step 1: First input, run untrained process.
    # Step 2: Second input, run untrained process, learns for the first time.
    # Step 3: It already learned, and there's a new input, so the trained process should run.
    # Step 4: Run trained process again.

    # The second expert:
    # Step 1: First input, run untrained process.
    # Step 2: Same input, nothing runs.
    # Step 3: Second input, run untrained process, learns for the first time.
    # Step 4: It already learned, and there's a new input, so the trained process should run.

    # indices of trained experts that are ran in a given step
    expected_indices_trained_seq = [
        torch.tensor([[]], dtype=torch.long, device=device).permute(1, 0),
        # the expected shape of this is [0,1], not [1,0]
        torch.tensor([[]], dtype=torch.long, device=device).permute(1, 0),
        torch.tensor([[0]], dtype=torch.long, device=device),
        torch.tensor([[0], [1]], dtype=torch.long, device=device),
    ]

    # indices of untrained experts that are ran in a givne step
    expected_indices_untrained_seq = [
        torch.tensor([[0], [1]], dtype=torch.long, device=device),
        torch.tensor([[0]], dtype=torch.long, device=device),
        torch.tensor([[1]], dtype=torch.long, device=device),
        torch.tensor([[]], dtype=torch.long, device=device).permute(1, 0),
    ]

    def _last_called_args(mock):
        # this extracts parameters that are passed to the create method
        return [c for c in mock.mock_calls if c[0] == ''][-1][1]

    with patch.object(TrainedForwardProcessFactory,
                      'create',
                      wraps=trained_factory.create) as trained_mock:
        with patch.object(UntrainedForwardProcessFactory,
                          'create',
                          wraps=untrained_factory.create) as untrained_mock:
            pooler = create_tp_flock(flock_size=flock_size,
                                     n_cluster_centers=n_cluster_centers,
                                     batch_size=batch_size,
                                     buffer_size=buffer_size,
                                     seq_length=seq_length,
                                     trained_forward_factory=trained_factory,
                                     untrained_forward_factory=untrained_factory,
                                     device=device)

            trained_call_indices_seq = []
            untrained_call_indices_seq = []
            # for each sequence of input clusters
            for input_clusters in input_clusters_seq:
                pooler.forward_learn(input_clusters)

                # read the parameters passed to the create method
                last_trained_create_call = _last_called_args(trained_mock)
                last_untrained_create_call = _last_called_args(untrained_mock)

                # store just the indices with which the process was created.
                trained_call_indices_seq.append(last_trained_create_call[4])
                untrained_call_indices_seq.append(last_untrained_create_call[4])

    # test that the forward_learn process correctly determined the trained and untrained indices
    assert same(trained_call_indices_seq[0], expected_indices_trained_seq[0])
    assert same(untrained_call_indices_seq[0], expected_indices_untrained_seq[0])
    assert same(trained_call_indices_seq[1], expected_indices_trained_seq[1])
    assert same(untrained_call_indices_seq[1], expected_indices_untrained_seq[1])
    assert same(trained_call_indices_seq[2], expected_indices_trained_seq[2])
    assert same(untrained_call_indices_seq[2], expected_indices_untrained_seq[2])
    assert same(trained_call_indices_seq[3], expected_indices_trained_seq[3])
    assert same(untrained_call_indices_seq[3], expected_indices_untrained_seq[3])

class TestTPFlock:
    @pytest.mark.parametrize('device', ['cpu', 'cuda'])
    @pytest.mark.parametrize('data, seqs, likelihoods, n_top_sequences, expected_output', [
        (
                [[0, 0, 1, 0]],
                [[[0, 2, 0], [0, 2, 1]]],
                [[1.0, 0.5]],
                1,
                [[0.666, 0.0, 0.333, 0.0]]
        ),
        (
                [[0, 0, 1, 0]],
                [[[0, 2, 0], [0, 2, 1]]],
                [[0.5, 1.0]],
                1,
                [[0.333, 0.333, 0.333, 0.0]]
        ),
        (
                [[0, 0, 1, 0]],
                [[[0, 2, 0], [0, 2, 1], [0, 1, 3]]],
                [[0.5, 1.0, 0.5]],
                2,
                [[0.5000, 0.1667, 0.3333, 0.0000]]
        )
    ])
    def test_inverse_projection(self, data, seqs, likelihoods, n_top_sequences, expected_output, device):
        float_type = get_float(device)
        t_data = torch.tensor(data, dtype=float_type, device=device)
        t_seqs = torch.tensor(seqs, dtype=torch.int64, device=device)
        t_likelihoods = torch.tensor(likelihoods, dtype=float_type, device=device)
        t_expected_output = torch.tensor(expected_output, dtype=float_type, device=device)
        flock_size, n_frequent_seq, seq_length = t_seqs.shape
        n_cluster_centers = t_data.shape[-1]

        flock = create_tp_flock(flock_size=flock_size, seq_length=seq_length, seq_lookahead=1,
                                n_frequent_seq=n_frequent_seq, n_cluster_centers=n_cluster_centers, device=device)

        flock.frequent_seqs = t_seqs
        flock.frequent_seq_likelihoods_priors_clusters_context = t_likelihoods
        output = flock.inverse_projection(t_data, n_top_sequences=n_top_sequences)
        assert same(t_expected_output, output, eps=1e-3)


# Say you have an array for which the ith element is the price of a given stock on day i.
#
# Design an algorithm to find the maximum profit. You may complete at most k transactions.
#
# Note:
# You may not engage in multiple transactions at the same time (ie, you must sell the stock before you buy again).

class Solution(object):
    def maxProfit(self, k, prices):
        """
        :type k: int
        :type prices: List[int]
        :rtype: int
        """
        n = len(prices)
        if n < 2:
            return 0
        # k is big enougth to cover all ramps.
        if k >= n / 2:
            return sum(i - j for i, j in zip(prices[1:], prices[:-1]) if i - j > 0)
        globalMax = [[0] * n for _ in range(k + 1)]
        for i in range(1, k + 1):
            # The max profit with i transations and selling stock on day j.
            localMax = [0] * n
            for j in range(1, n):
                profit = prices[j] - prices[j - 1]
                localMax[j] = max(
                        # We have made max profit with (i - 1) transations in
                        # (j - 1) days.
                        # For the last transation, we buy stock on day (j - 1)
                        # and sell it on day j.
                        globalMax[i - 1][j - 1] + profit,
                        # We have made max profit with (i - 1) transations in
                        # (j - 1) days.
                        # For the last transation, we buy stock on day j and
                        # sell it on the same day, so we have 0 profit, apparently
                        # we do not have to add it.
                        globalMax[i - 1][j - 1],  # + 0,
                        # We have made profit in (j - 1) days.
                        # We want to cancel the day (j - 1) sale and sell it on
                        # day j.
                        localMax[j - 1] + profit)
                globalMax[i][j] = max(globalMax[i][j - 1], localMax[j])
        return globalMax[k][-1]

Python/188. Best Time to Buy and Sell Stock IV.py


import csv
from flask import Flask,render_template,request,redirect

app = Flask(__name__)

students=[]

@app.route("/")
def home_page():
    return render_template("home.html")

@app.route("/register",methods=['POST'])
def register():

    if not request.form.get("name") or not request.form.get('dorm') :
       return render_template("failure.html")

    file =open("registered.csv","a")
    writer=csv.writer(file)
    writer.writerow((request.form.get("name"),request.form.get('dorm')))
    file.close()

    return render_template("success.html")

@app.route("/registered")
def registered():

    with open("registered.csv","r") as file:
         reader = csv.reader(file)
         students = list(reader)
    return render_template("registered.html",students=students)


import falcon
import json
import pandas as pd
import pickle
import io,sys
import pandas.core.indexes
import numpy as np
sys.modules['pandas.indexes'] = pandas.core.indexes
from sklearn.preprocessing import LabelEncoder,OneHotEncoder
import lime
import lime.lime_tabular

class Data():
    def __init__(self):
        self.x_train = ""
        self.x_test = ""
        self.y_test = ""
        self.inversed_x_test = ""
        self.categorical_features = ""
        self.categorical_names = ""
        self.feature_names = ""
        self.label_names = ""
        self.trained_model = ""
        self.accessor = ""
        self.isOneHot = False
        self.prediction_result = ""
        self.oneHotEncoder = ""

    def input_data(self,target_name,data):
        if target_name == "X_train":
            self.x_train = data
        elif target_name == "X_test":
            self.x_test = data
        elif target_name == "y_test":
            self.y_test = data
        elif target_name == "Categorical Features":
            self.categorical_features = data
        elif target_name == "Categorical Names":
            self.categorical_names = data
        elif target_name == "Feature Names":
            self.feature_names = data
        elif target_name == "Label Names":
            self.label_names = data
        elif target_name == "Trained Model":
            self.trained_model = data

    def inverse_and_add_prediction(self):
        test = self.x_test.copy()
        ## categorical features inverse
        categorical_names = self.categorical_names
        categorical_features = self.categorical_features
        for i in categorical_features:
            le = LabelEncoder()
            le.fit(categorical_names[i])
            test.iloc[:,i] = le.inverse_transform(test.iloc[:,i])
        ## add prediction_result
        test[len(test.columns)] = self.prediction_result
        ## add index
        test[len(test.columns)] = list(np.arange(len(test)))
        test.columns = self.accessor
        self.inversed_x_test = test

    def create_accessor(self):
        self.feature_names = np.append(self.feature_names,"Prediction")
        self.feature_names = np.append(self.feature_names,"Index")
        feature_names = self.feature_names
        res = []
        for i in feature_names:
            res.append(i.replace(" ","").lower())
        self.accessor = res

    def addPrediction(self,isOneHot):
        self.isOneHot = isOneHot
        model = self.trained_model
        if not isinstance(self.x_train,np.ndarray):
            train = self.x_train.as_matrix()
        else:
            train = self.x_train
        if not isinstance(self.x_test,np.ndarray):
            test = self.x_test.as_matrix()
        else:
            test = self.x_test
        if isOneHot:
            encoder = OneHotEncoder(categorical_features=self.categorical_features)
            alldata = np.concatenate([train,test])
            encoder.fit(alldata)
            self.oneHotEncoder = encoder
            encoded_test = encoder.transform(test)
            prediction_result = model.predict_proba(encoded_test)
        else:
            ## for example, LightGBM
            prediction_result = model.predict_proba(test)
        t_res = []
        for i in prediction_result:
            t_res.append(i[0])
        self.prediction_result = t_res

class TabularLIME():
    def __init__(self):
        self.explainer = ""
        self.predict_fn = ""
        self.exp_instance = ""
        self.expScore = ""
        self.expList = ""
        self.expProba = ""

    def trainLIME(self,data):
        if not isinstance(data.x_train,np.ndarray):
            train = data.x_train.as_matrix()
        else:
            train = data.x_train
        self.explainer = lime.lime_tabular.LimeTabularExplainer(train,
        feature_names=data.feature_names, class_names=data.label_names,
        categorical_features=list(data.categorical_features),
        categorical_names=data.categorical_names, kernel_width=3)
        ## custom predict function
        self.predict_fn = lambda x: data.trained_model.predict_proba(data.oneHotEncoder.transform(x)).astype(float)
        print("train tabular LIME complete")

    def doExplanation(self,index,data):
        test = data.x_test.as_matrix()
        exp = self.explainer.explain_instance(test[int(index)], self.predict_fn, num_features=len(data.feature_names))
        self.exp_instance = exp
        self.expScore = exp.score
        self.expProba = exp.predict_proba
        self.expList = exp.as_list()

def check_train_test_dtype(df):
    message = "Data check passed"
    result = True
    if len(df.columns) < 1:
        message = "This dataframe is empty"
        result = False
        return message,result
    for column in df:
        if df[column].dtype not in ["int32","int64","float32","float64","int","float"]:
            message = "This data contains something other than int or float column"
            result = False
    return message,result

def check_target_test_dtype(target_test):
    message = "Data check passed"
    result = True
    if len(target_test.columns) > 1:
        message = "This target data is over 1 column"
        result = False
    if target_test.iloc[:,0].dtype not in ["int32","int64","int"]:
        message = "This data contains something other than int column"
        result = False

    return message,result

def check_pickle_list_dtype(pickle_list,target_name):
    message = "Data check passed"
    result = True
    ## nested array or not
    if isinstance(pickle_list[0],list):
        message = "Nested array can not be used"
        result = False
    if target_name == "Feature Names":
        for i in pickle_list:
            if not isinstance(i,str):
                message = "This list contains not string data"
                result = False
    elif target_name == "Categorical Features":
        for i in pickle_list:
            if not isinstance(i,int) and not isinstance(i,np.int32) and not isinstance(i,np.int64):
                message = "This list contains not integer data"
                result = False

    return message,result

def check_pickle_dict_dtype(pickle_dict):
    message = "Data check passed"
    result = True
    ## dict or not
    if not isinstance(pickle_dict,dict):
        message = "Not dict object can not be used"
        result = False
    return message,result

def methods_for_pickle(message,result,df,file,target_name):
    pickle_file = pickle.loads(file)
    ## if dataframe
    if isinstance(pickle_file,pd.core.frame.DataFrame):
        df = pickle_file
        if target_name == "X_train" or target_name == "X_test":
            message,result = check_train_test_dtype(df)
        elif target_name == "y_test":
            message,result = check_target_test_dtype(df)
    ## if numpy.ndarray, convert dataframe
    elif isinstance(pickle_file,np.ndarray):
        if target_name == "X_train" or target_name == "X_test":
            df = pd.DataFrame(pickle_file)
            message,result = check_train_test_dtype(df)
        elif target_name == "y_test":
            df = pd.DataFrame(pickle_file)
            message,result = check_target_test_dtype(df)
        elif target_name == "Categorical Features" or target_name == "Feature Names" or target_name == "Label Names":
            message,result = check_pickle_list_dtype(pickle_file,target_name)
    elif isinstance(pickle_file,list):
        if target_name == "Categorical Features" or target_name == "Feature Names" or target_name == "Label Names":
            np_pickle_file = np.array(pickle_file)
            df = np_pickle_file
            message,result = check_pickle_list_dtype(np_pickle_file,target_name)
        else:
            message = "This data is not pickled list although this data should be list"
            result = False
    elif isinstance(pickle_file,dict):
        df = pickle_file
        if target_name == "Categorical Names":
            message,result = check_pickle_dict_dtype(pickle_file)
        else:
            message = "This data is not pickled dict although this data should be dict"
            result = False
    else:
        message = "This data is not dataframe,numpy.ndarray,list or dict"
        result = False
    if target_name == "Trained Model":
        message = "Data Check Passed"
        df = pickle_file
        result = True
    return message,result,df

def check_data_content(file,target_name,file_type):
    message = "Data check passed"
    result = True
    df = ""
    if target_name == "X_train" or target_name == "X_test":
        if file_type == ".csv":
            df = pd.read_csv(io.BytesIO(file))
            ## df can contain number only
            message,result = check_train_test_dtype(df)
        elif file_type == ".pickle":
            message,result,df = methods_for_pickle(message,result,df,file,target_name)
        else:
            message = "This data is not dataframe or numpy.ndarray"
            result = False

    elif target_name == "y_test":
        if file_type == ".csv":
            df = pd.read_csv(io.BytesIO(file))
            message,result = check_target_test_dtype(df)
        elif file_type == ".pickle":
            message,result,df = methods_for_pickle(message,result,df,file,target_name)

    elif target_name == "Categorical Features" or target_name == "Feature Names" or target_name == "Categorical Names" or target_name == "Label Names":
        ## list only(pickle)
        if file_type == ".pickle":
            message,result,df = methods_for_pickle(message,result,df,file,target_name)
        else:
            message = "This data not pickle"
            result = False
            df = ""

    elif target_name == "Trained Model":
        ## pickle,h5 model file
        if file_type == ".pickle":
            message,result,df = methods_for_pickle(message,result,df,file,target_name)
        # elif file_type == ".h5":
        #     message,result,df = methods_for_h5(message,result,df,file,target_name)
        # elif file_type == ".model":
        #     message,result,df = methods_for_model(message,result,df,file,target_name)
        else:
            message = "This data not pickle"
            result = False
            df = ""

    return message,result,df


import cv2
import numpy as np

fld = cv2.ximgproc.createFastLineDetector()
"""
The array is colored
"""
def get_otsu_threshold(img:np.ndarray):
    gray_img= cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    return cv2.threshold(gray_img,100,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)[1]

def get_kmeans_on_image(img:np.ndarray):
    datapoints= img.reshape((-1,3)).astype(np.float32)
    # print(datapoints.shape)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    ret,labels,centers= cv2.kmeans(datapoints,3, None,criteria,10,cv2.KMEANS_RANDOM_CENTERS)
    labels=labels.reshape((img.shape[0],img.shape[1]))
    # print(labels.shape)
    images= []
    for i in range(3):
        output= img.copy()
        # print(output.shape)
        output[labels!=i]=np.uint8([0,0,0])
        images.append(output)
    return images

def detect_lines(img:np.ndarray):
    gray_img= cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    gray_img= cv2.threshold(gray_img,30,255,cv2.THRESH_BINARY)[1]
    blurred= cv2.GaussianBlur(gray_img,(5,3),33)
    lines= fld.detect(blurred, 150)
    output_img= img.copy()
    if lines is None:return
    for line in lines:
        line= line[0].astype('int')
        cv2.line(output_img,(line[0],line[1]),(line[2],line[3]),(255,0,0),5)
    cv2.imshow('Line Detector', output_img)

def draw_contours(img:np.ndarray):
    processed= cv2.medianBlur(img,21)
    processed= get_otsu_threshold(img)
    processed= cv2.Canny(img,30,130)
    cv2.imshow('Canny',processed)
    contours, hierarchy = cv2.findContours(processed,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    output_img= np.zeros_like(img,dtype=np.uint8)
    cv2.drawContours(output_img, contours, -1, (0, 255, 0), 3)

    cv2.imshow('Contours', output_img)

cap= cv2.VideoCapture(0)

while True:
    frame= cap.read()[1]

    cv2.imshow('Original',frame)
    detect_lines(frame)
    draw_contours(frame)
    if cv2.waitKey(1) == 27:
        break
cv2.destroyAllWindows()
cap.release()


class Solution:
    def find132pattern(self, nums: List[int]) -> bool:
        te=float(-inf)
        stack=[]
        for i in range(len(nums)-1,-1,-1):
            if nums[i]<te:
                return True
            while stack and stack[-1]<nums[i]:
                te=stack.pop()

            stack.append(nums[i])
        return False

# class Solution:

#     def find132pattern(self, nums):
#         stack, s3 = [], -float("inf")
#         for n in nums[::-1]:
#             if n < s3: return True
#             while stack and stack[-1] < n: s3 = stack.pop()
#             stack.append(n)
#         return False

# #     My understanding:
# # You go from the right side, and find the first element that is bigger than one or more elements that are further to the right. If that happens, you want to find the biggest number that is smaller than the current element. You do that by popping out all of those smaller elements.

# # Why you want to find the biggest number that is smaller than the current element? Because you want to maximize the chance that you find something smaller than the number going left. If you find such smaller one, then return True.
# # Note that the elements in the stack are stored in descending order, all elements that are append to the stack are smaller than what's already in the stack. This is important, since we want to guarantee that the last element that we pop out of the stack is the biggest value that is smaller than the current value.


import argparse
import benchmarks
import os
import planning
import sys
import json
import networkx as nx
import matplotlib.pyplot as plt
from benchmarks import *
from networkx.drawing.nx_agraph import graphviz_layout
from moz_sql_parser import parse
from enum import Enum
from planning import *

def load_schema(schema_path):
    base_tables = {}
    f = open(schema_path)
    lines = f.readlines()
    most_recent_table = None
    for line in lines:
        if 'CREATE' in line:
            table_name = line.split('TABLE')[1].replace('(', '').replace('`', '').replace('\n', '')
            most_recent_table = table_name.strip()
            base_tables[most_recent_table] = []
        elif len(line.strip()) > 0 and line.strip()[0] == '`':
            line = line.split()
            col_name = line[0].replace('`', '')
            base_tables[most_recent_table].append(col_name)
        else:
            continue
    return base_tables

def load_policies(schema, benchmark):
    if benchmark == 'hotcrp':
        event_chain = hotcrp_policy_nodes
    elif benchmark == 'twitter':
        event_chain = twitter_policy_nodes
    else:
        raise NotImplementedError

    name_to_op = {}
    for op in event_chain:
        name_to_op[op.new_view_name] = op

    starts = []
    chains = []
    for node in event_chain:
        if node.exported_as is not None:
            starts.append(node)

    for i, node in enumerate(starts):
        chain = [node]
        parents = node.tables
        while len(parents) > 0:
            new_parents = []
            for parent in parents:
                if parent in name_to_op:
                    parent_node = name_to_op[parent]
                    chain.append(parent_node)
                    for tbl in parent_node.tables:
                        new_parents.append(tbl)
            parents = new_parents
        chains.append(chain)

    policies = []
    for chain in chains:
        chain.reverse()
        policy = Function(chain, schema)
        policy = policy.to_dataflow(schema)
        policies.append(policy)
    return policies

def load_queries(schema, benchmark):
    if benchmark == 'hotcrp':
        full_query = hotcrp_query_nodes
    elif benchmark == 'twitter':
        full_query = twitter_full_query
    else:
        raise NotImplementedError

    full_query = Function(full_query, schema)
    full_query = full_query.to_dataflow(schema)
    print("FULL QUERY: {}".format(full_query))
    return [full_query]

def visualize(graph):
    G = nx.DiGraph()

    for node in graph.keys():
        if node.operation_type == None:
            G.add_node(node.name)

    for node, edges in graph.items():
        for edge in edges:
            G.add_node(edge.name)
            G.add_edge(node.name, edge.name)

    pos = graphviz_layout(G, prog='dot')
    nx.draw(G, pos, with_labels=True, arrows=True)

    plt.show()

def main():
    parser = argparse.ArgumentParser(description='Select benchmark.')
    parser.add_argument('--benchmark', type=str, default='hotcrp')
    args = parser.parse_args()

    if args.benchmark == 'hotcrp':
        schema_path = '../benchmarks/hotcrp/schema.sql'
    elif args.benchmark == 'twitter':
        schema_path = '../benchmarks/twitter/twitter-schema.sql'
    else:
        raise NotImplementedError

    schema = load_schema(schema_path)
    queries = load_queries(schema, args.benchmark)
    # visualize(queries[0])
    policies = load_policies(schema, args.benchmark)

    for i, policy in enumerate(policies):
        print("POLICY {}: {}".format(i, policy))
        visualize(policy)
    final_graph = planning(queries, policies)
    print('final graph: {}'.format(final_graph))
    # visualize(final_graph)
    for graph in final_graph:
        visualize(graph)

if __name__ == '__main__':
    main()


import json

def AddNameBytes(DataFrame, str):
    for char in str:
        DataFrame.append(ord(char))
    return DataFrame

# ACK byte = 0x00: reserved for future use
# CMDNUMBER (2 bytes): 0x00 0xxx

def Error_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    # L     ID      CMD     FT      CMDNUMBER       ACK
    DataFrame = [0x06, ID, 0x01, 0x02, 0x00, CmdNumber, 0x00]
    return DataFrame

def Event_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    # L     ID      CMD     FT      CMNNUMBER       REMAINING SHOT (if exist)   ACK
    if "RemainingShot" in temp:
        RemainingShot = int(temp["RemainingShot"])
        DataFrame = [0x07, ID, 0x02, 0x02, 0x00, CmdNumber, RemainingShot, 0x00]
    else:
        DataFrame = [0x06, ID, 0x02, 0x02, 0x00, CmdNumber, 0x00]
    return DataFrame

def Card_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    if (temp["AuthenticationStatus"] == "Authenticated"):
        AuthenticationStatus = 0x00
    else:
        AuthenticationStatus = 0x01

    if (temp["CardHolderRole"] == "MANAGER"):
        CardHolderRole = 0x02
    elif (temp["CardHolderRole"] == "BARTENDER"):
        CardHolderRole = 0x01
    else:
        CardHolderRole = 0x00

    CardHolderName = temp["CardHolderName"]
    NameBytes = CardHolderName.encode('utf-8')
    NameLength = len(NameBytes)

    FrameLength = 8 + NameLength
    # L     ID      CMD     FT      CMDNUMBER   ACK       AuStatus        Role        Name
    DataFrame = [FrameLength, ID, 0x03, 0x02, 0x00, CmdNumber, 0x00, AuthenticationStatus, CardHolderRole]
    DataFrame = AddNameBytes(DataFrame, CardHolderName)
    return DataFrame

def Barcode_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    if (temp["AuthenticationStatus"] == "Authenticated"):
        AuthenticationStatus = 0x00
    else:
        AuthenticationStatus = 0x01
    RemainingShot = int(temp["RemainingShot"])
    Compare = int(temp["Compare"])
    SpiritName = temp["SpiritName"]
    NameBytes = SpiritName.encode('utf-8')
    NameLength = len(NameBytes)

    FrameLength = 8 + NameLength
    # L     ID      CMD     FT      CMDNUMBER   ACK       AuStatus        Compare        RemainingShot      (no name length)      SpiritName
    DataFrame = [FrameLength, ID, 0x04, 0x02, 0x00, CmdNumber, 0x00, AuthenticationStatus, Compare, RemainingShot]
    DataFrame = AddNameBytes(DataFrame, SpiritName)
    return DataFrame

def DispensingSpirit_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    SessionTotalShot = int(temp["SessionTotalShot"])
    RemainingShot = int(temp["RemainingShot"])
    # L     ID      CMD     FT      CMDNUMBER   ACK       SessionTotalShot        RemainingShot
    DataFrame = [0x08, ID, 0x05, 0x02, 0x00, CmdNumber, 0x00, SessionTotalShot, RemainingShot]
    return DataFrame

def Synchronize_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    #DispenserID = int(temp["DispenserID"])
    RemainingShot = int(temp["RemainingShot"])

    SpiritName = temp["SpiritName"]
    NameBytes = SpiritName.encode('utf-8')
    NameLength = len(NameBytes)

    FrameLength = 7 + NameLength
    # L     ID      CMD     FT      CMDNUMBER   ACK         RemainingShot       SpiritName
    DataFrame = [FrameLength, ID, 0x06, 0x02, 0x00, CmdNumber, 0x00, RemainingShot]
    DataFrame = AddNameBytes(DataFrame, SpiritName)
    return DataFrame

def UpdateStatus_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    # L     ID      CMD     FT      CMDNUMBER   ACK
    DataFrame = [0x06, ID, 0x07, 0x02, 0x00, CmdNumber, 0x00]
    return DataFrame

def CheckConnection_FrameEncode(Payload):
    temp = json.loads(Payload)      # parse JSON payload
    ID = int(temp["ID"])
    CmdNumber = int(temp["CmdNumber"])
    # L     ID      CMD     FT      CMDNUMBER   ACK
    DataFrame = [0x06, ID, 0x08, 0x02, 0x00, CmdNumber, 0x00]
    return DataFrame


#Chapter 6 Exercises
#1. Pizza Toppings ####################################################
print("Exercise 1: Pizza Toppings")

topping = "\nWhat toppings do you like on your pizza?"
topping += "\nEnter 'quit' when you are done inputing all of it: "

while True:
    topping = input(topping)
    if topping != 'quit': # If user types quit the code ends
        print(f"  The {topping} is added on the pizza.")
    else:
        break

#2. Movie Tickets ####################################################
print("Exercise 2: Movie Tickets")

age = "How old are you?"
age += "\nEnter 'done' when you are finished. "

while True:
    age = input(age)
    if age == 'done': # This code will finish the code if the user types done
        break
    age = int(age)

    if age < 3: # If the age is lower than 3
        print(" Toddlers get to watch the movie for free")
    elif age < 13:# If the age is lower than 13
        print("Children who are under 13 years old get to watch the movies for 10$")
    else:
        print("Your ticket is $15") # If not any of the others

#3. Infinity ####################################################
print("Exercise 3: Infinity")

while True:
    print("hello world")
    break #Remove the break statement to stop the infinite loop

#4. Deli ####################################################
print("Exercise 4: Deli")

sandwich_orders = ['Bacon', 'BLT', 'Egg', 'Cheese']
sandwich_done = []

while sandwich_orders:
    current_sandwich = sandwich_orders.pop()
    print(f"I'm working on your {current_sandwich} sandwich.\n")
    sandwich_done.append(current_sandwich)

for sandwich in sandwich_done:
    print(f"I made a {sandwich} sandwich.")

#5. No Pastrami ####################################################
print("Exercise 5: No Pastrami")

sandwich_orders = [
    'pastrami', 'Bacon', 'BLT', 'Egg', 'Cheese']
sandwich_done = []

print("I'm sorry, we're all out of pastrami today.")
while 'pastrami' in sandwich_orders:
    sandwich_orders.remove('pastrami') # Removes Pastrami

print("\n")
while sandwich_orders:
    current_sandwich = sandwich_orders.pop()
    print(f"I'm working on your {current_sandwich} sandwich.") #Prints all the codes
    sandwich_done.append(current_sandwich)

print("\n")
for sandwich in sandwich_done:
    print(f"I made a {sandwich} sandwich.")


import sys

if __name__ == '__main__':
	nums = [1, 1]

	for i in range(2, 51):
		nums.append(nums[i-1] + nums[i-2])

	while True:
		x = int(sys.stdin.readline())

		if x == 0:
			break

		print(nums[x])


import time
import sys as eeeeee
from webbrowser import open_new_tab as eeeeeeeeeeeeeeeeee
from plyer import notification as eeee
from typing import List as eee
from base64 import urlsafe_b64decode as e
from traceback import format_exception as eeeeeeeeeeee

def not_my_toes(*args, **kwargs):
    print(*args, **kwargs)

def main(args: eee[str] = None):
    if args is not None or len(args) > 0:
        not_my_toes(e(b'T2ggeW91IHBhc3NlZCBzb21lIGFyZ3VtZW50cywgYnV0IHNhZGx5IGlkZ2FmLiBSaXAgYm96bw==').decode("utf-8"))

    eeeeeeeeeeeeeeeeee(e(b'aHR0cHM6Ly93d3cueW91dHViZS5jb20vd2F0Y2g_dj1kUXc0dzlXZ1hjUQ==').decode("utf-8"))
    while True:
        with open("NewFile1.txt", "r") as file:
            try:
                for eeeee in file:
                    eeee.notify(
                        title="Garbage",
                        message=eeeee,
                        timeout=3,
                    )
                    time.sleep(5)
            except BaseException as eeeeee:
                eeee.notify(
                    title="Garbage",
                    message=f"Uh oh, cringe: {''.join(eeeeeeeeeeee(eeeeee))}",
                    timeout=3
                )
        not_my_toes("A")

if __name__ == "__main__":
    main(eeeeee.argv)


import pytest
import torch

from inseq.utils.misc import pretty_tensor
from inseq.utils.torch_utils import filter_logits

@pytest.mark.parametrize(
    ("tensor", "output"),
    [
        (None, "None"),
        (torch.tensor([1, 1], dtype=torch.long), "torch.int64 tensor of shape [2]"),
        (
            torch.tensor([[1, 1, 1], [1, 1, 1]], dtype=torch.long),
            "torch.int64 tensor of shape [2, 3]",
        ),
        (torch.randn(4, 1, 1, 10), "torch.float32 tensor of shape [4, 1, 1, 10]"),
        (torch.randn(2, 21, 1), "torch.float32 tensor of shape [2, 21, 1]"),
    ],
)
def test_pretty_tensor(tensor: torch.Tensor, output: str) -> None:
    assert pretty_tensor(tensor).startswith(output)

def test_probits2prob():
    # Test with batch of size > 1
    probits = torch.stack(
        [
            torch.arange(0, 30000, 1.0),
            torch.arange(30000, 60000, 1.0),
            torch.arange(60000, 90000, 1.0),
            torch.arange(90000, 120000, 1.0),
        ]
    )
    target_ids = torch.tensor([10, 77, 999, 1765]).unsqueeze(-1)
    probs = torch.gather(probits, -1, target_ids.T)
    assert probs.shape == (1, 4)
    assert torch.eq(probs, torch.tensor([10.0, 77.0, 999.0, 1765.0])).all()

    # Test with batch of size 1
    probits = torch.stack(
        [
            torch.arange(0, 30000, 1.0),
        ]
    )
    target_ids = torch.tensor([23456]).unsqueeze(-1)
    probs = torch.gather(probits, -1, target_ids.T)
    assert probs.shape == (1, 1)
    assert torch.eq(probs, torch.tensor([23456.0])).all()

def test_filter_logits():
    original_logits = torch.tensor(
        [
            [1.0, 2.0, 3.0, 4.0, 5.0],
            [2.0, 3.0, 4.0, 5.0, 6.0],
            [3.0, 4.0, 5.0, 6.0, 7.0],
            [4.0, 5.0, 6.0, 7.0, 8.0],
        ]
    )
    filtered_logits = filter_logits(original_logits, top_k=2)
    topk2 = original_logits.clone().index_fill(1, torch.tensor([0, 1, 2]), float("-inf"))
    assert torch.eq(filtered_logits, topk2).all()
    filtered_logits = filter_logits(original_logits, top_p=0.9)
    topp90 = original_logits.clone().index_fill(1, torch.tensor([0, 1]), float("-inf"))
    assert torch.eq(filtered_logits, topp90).all()
    contrast_logits = torch.tensor(
        [
            [15.0, 13.0, 11.0, 9.0, 7.0],
            [13.0, 11.0, 9.0, 7.0, 5.0],
            [11.0, 9.0, 7.0, 5.0, 3.0],
            [9.0, 7.0, 5.0, 3.0, 1.0],
        ]
    )
    filtered_logits, contrast_logits = filter_logits(original_logits, contrast_logits=contrast_logits, top_k=2)
    top2merged = original_logits.clone().index_fill(1, torch.tensor([2, 3, 4]), float("-inf"))
    assert torch.eq(filtered_logits, top2merged).all()