#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
config_test.py

Created by Chandrasekhar Ramakrishnan on 2017-02-10.
Copyright (c) 2017 Chandrasekhar Ramakrishnan. All rights reserved.
"""
import json
import os
import shutil

from . import config

def test_config_location_resolver():
    loc = config.ConfigLocation(['global'], 'user_home', '.obis')
    location_resolver = config.LocationResolver()
    assert location_resolver.resolve_location(loc) == os.path.join(os.path.expanduser("~"), '.obis')

def user_config_test_data_path():
    return os.path.join(os.path.dirname(__file__), '..', 'test-data', 'user_config')

def copy_user_config_test_data(tmpdir):
    config_test_data_src = user_config_test_data_path()
    config_test_data_dst = str(tmpdir.join(os.path.basename(config_test_data_src)))
    shutil.copytree(config_test_data_src, config_test_data_dst)
    return config_test_data_src, config_test_data_dst

def configure_resolver_for_test(resolver, tmpdir):
    resolver.set_resolver_location_roots('user_home', os.path.join(str(tmpdir), 'user_config'))

def test_read_config(tmpdir):
    copy_user_config_test_data(tmpdir)
    resolver = config.SettingsResolver().config
    configure_resolver_for_test(resolver, tmpdir)
    config_dict = resolver.config_dict()
    assert config_dict is not None
    with open(os.path.join(user_config_test_data_path(), ".obis", "config.json")) as f:
        expected_dict = json.load(f)
    assert config_dict['user'] == expected_dict['user']
    assert './.obis/config.json' == resolver.local_public_properties_path()

def test_write_config(tmpdir):
    copy_user_config_test_data(tmpdir)
    resolver = config.SettingsResolver().config
    configure_resolver_for_test(resolver, tmpdir)
    config_dict = resolver.config_dict()
    assert config_dict is not None
    with open(os.path.join(user_config_test_data_path(), ".obis", "config.json")) as f:
        expected_dict = json.load(f)
    assert config_dict['openbis_url'] == expected_dict['openbis_url']
    assert config_dict['user'] == expected_dict['user']

    resolver.set_value_for_parameter('user', 'new_user', 'local')
    config_dict = resolver.config_dict()
    assert config_dict['openbis_url'] == expected_dict['openbis_url']
    assert config_dict['user'] == 'new_user'


# coding=utf-8
# Copyright 2018-2023 EvaDB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
from pathlib import Path

from evadb.catalog.catalog_type import ColumnType, NdArrayType, VectorStoreType
from evadb.expression.abstract_expression import ExpressionType
from evadb.expression.comparison_expression import ComparisonExpression
from evadb.expression.constant_value_expression import ConstantValueExpression
from evadb.expression.function_expression import FunctionExpression
from evadb.expression.tuple_value_expression import TupleValueExpression
from evadb.parser.alias import Alias
from evadb.parser.create_index_statement import CreateIndexStatement
from evadb.parser.create_mat_view_statement import CreateMaterializedViewStatement
from evadb.parser.create_statement import (
    ColConstraintInfo,
    ColumnDefinition,
    CreateTableStatement,
)
from evadb.parser.create_udf_statement import CreateUDFStatement
from evadb.parser.delete_statement import DeleteTableStatement
from evadb.parser.drop_object_statement import DropObjectStatement
from evadb.parser.insert_statement import InsertTableStatement
from evadb.parser.load_statement import LoadDataStatement
from evadb.parser.parser import Parser
from evadb.parser.rename_statement import RenameTableStatement
from evadb.parser.select_statement import SelectStatement
from evadb.parser.statement import AbstractStatement, StatementType
from evadb.parser.table_ref import JoinNode, TableInfo, TableRef, TableValuedExpression
from evadb.parser.types import (
    FileFormatType,
    JoinType,
    ObjectType,
    ParserOrderBySortType,
)

class ParserTests(unittest.TestCase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def test_create_index_statement(self):
        parser = Parser()

        create_index_query = "CREATE INDEX testindex ON MyVideo (featCol) USING FAISS;"
        evadb_stmt_list = parser.parse(create_index_query)

        # check stmt itself
        self.assertIsInstance(evadb_stmt_list, list)
        self.assertEqual(len(evadb_stmt_list), 1)
        self.assertEqual(evadb_stmt_list[0].stmt_type, StatementType.CREATE_INDEX)

        expected_stmt = CreateIndexStatement(
            "testindex",
            TableRef(TableInfo("MyVideo")),
            [
                ColumnDefinition("featCol", None, None, None),
            ],
            VectorStoreType.FAISS,
        )
        actual_stmt = evadb_stmt_list[0]
        self.assertEqual(actual_stmt, expected_stmt)

        # create index on UDF expression
        create_index_query = (
            "CREATE INDEX testindex ON MyVideo (FeatureExtractor(featCol)) USING FAISS;"
        )
        evadb_stmt_list = parser.parse(create_index_query)

        # check stmt itself
        self.assertIsInstance(evadb_stmt_list, list)
        self.assertEqual(len(evadb_stmt_list), 1)
        self.assertEqual(evadb_stmt_list[0].stmt_type, StatementType.CREATE_INDEX)

        func_expr = FunctionExpression(None, "FeatureExtractor")
        func_expr.append_child(TupleValueExpression("featCol"))
        expected_stmt = CreateIndexStatement(
            "testindex",
            TableRef(TableInfo("MyVideo")),
            [
                ColumnDefinition("featCol", None, None, None),
            ],
            VectorStoreType.FAISS,
            func_expr,
        )
        actual_stmt = evadb_stmt_list[0]
        self.assertEqual(actual_stmt, expected_stmt)

    @unittest.skip("Skip parser exception handling testcase, moved to binder")
    def test_create_index_exception_statement(self):
        parser = Parser()

        create_index_query = (
            "CREATE INDEX testindex USING FAISS ON MyVideo (featCol1, featCol2);"
        )

        with self.assertRaises(Exception):
            parser.parse(create_index_query)

    def test_explain_dml_statement(self):
        parser = Parser()

        explain_query = "EXPLAIN SELECT CLASS FROM TAIPAI;"
        evadb_statement_list = parser.parse(explain_query)

        # check explain stmt itself
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.EXPLAIN)

        # check inner stmt
        inner_stmt = evadb_statement_list[0].explainable_stmt
        self.assertEqual(inner_stmt.stmt_type, StatementType.SELECT)

        # check inner stmt from
        self.assertIsNotNone(inner_stmt.from_table)
        self.assertIsInstance(inner_stmt.from_table, TableRef)
        self.assertEqual(inner_stmt.from_table.table.table_name, "TAIPAI")

    def test_explain_ddl_statement(self):
        parser = Parser()

        select_query = """SELECT id, Yolo(frame).labels FROM MyVideo
                        WHERE id<5; """
        explain_query = "EXPLAIN CREATE MATERIALIZED VIEW uadtrac_fastRCNN (id, labels) AS {}".format(
            select_query
        )

        evadb_statement_list = parser.parse(explain_query)

        # check explain stmt itself
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.EXPLAIN)

        # check inner stmt
        inner_stmt = evadb_statement_list[0].explainable_stmt
        self.assertEqual(inner_stmt.stmt_type, StatementType.CREATE_MATERIALIZED_VIEW)

        # check inner stmt from
        self.assertIsNotNone(
            inner_stmt.view_info, TableRef(TableInfo("uadetrac_fastRCNN"))
        )

    def test_create_table_statement(self):
        parser = Parser()

        single_queries = []
        single_queries.append(
            """CREATE TABLE IF NOT EXISTS Persons (
                  Frame_ID INTEGER UNIQUE,
                  Frame_Data TEXT(10),
                  Frame_Value FLOAT(1000, 201),
                  Frame_Array NDARRAY UINT8(5, 100, 2432, 4324, 100)
            );"""
        )

        expected_cci = ColConstraintInfo()
        expected_cci.nullable = True
        unique_cci = ColConstraintInfo()
        unique_cci.unique = True
        unique_cci.nullable = False
        expected_stmt = CreateTableStatement(
            TableInfo("Persons"),
            True,
            [
                ColumnDefinition("Frame_ID", ColumnType.INTEGER, None, (), unique_cci),
                ColumnDefinition(
                    "Frame_Data", ColumnType.TEXT, None, (10,), expected_cci
                ),
                ColumnDefinition(
                    "Frame_Value", ColumnType.FLOAT, None, (1000, 201), expected_cci
                ),
                ColumnDefinition(
                    "Frame_Array",
                    ColumnType.NDARRAY,
                    NdArrayType.UINT8,
                    (5, 100, 2432, 4324, 100),
                    expected_cci,
                ),
            ],
        )

        for query in single_queries:
            evadb_statement_list = parser.parse(query)
            self.assertIsInstance(evadb_statement_list, list)
            self.assertEqual(len(evadb_statement_list), 1)
            self.assertIsInstance(evadb_statement_list[0], AbstractStatement)
            self.assertEqual(evadb_statement_list[0], expected_stmt)

    def test_create_table_statement_with_rare_datatypes(self):
        parser = Parser()
        query = """CREATE TABLE IF NOT EXISTS Dummy (
                  C NDARRAY UINT8(5),
                  D NDARRAY INT16(5),
                  E NDARRAY INT32(5),
                  F NDARRAY INT64(5),
                  G NDARRAY UNICODE(5),
                  H NDARRAY BOOLEAN(5),
                  I NDARRAY FLOAT64(5),
                  J NDARRAY DECIMAL(5),
                  K NDARRAY DATETIME(5)
            );"""

        evadb_statement_list = parser.parse(query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertIsInstance(evadb_statement_list[0], AbstractStatement)

    def test_create_table_statement_without_proper_datatype(self):
        parser = Parser()
        query = """CREATE TABLE IF NOT EXISTS Dummy (
                  C NDARRAY INT(5)
                );"""

        with self.assertRaises(Exception):
            parser.parse(query)

    def test_create_table_exception_statement(self):
        parser = Parser()

        create_table_query = "CREATE TABLE ();"

        with self.assertRaises(Exception):
            parser.parse(create_table_query)

    def test_rename_table_statement(self):
        parser = Parser()
        rename_queries = "RENAME TABLE student TO student_info"
        expected_stmt = RenameTableStatement(
            TableRef(TableInfo("student")), TableInfo("student_info")
        )
        evadb_statement_list = parser.parse(rename_queries)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.RENAME)

        rename_stmt = evadb_statement_list[0]
        self.assertEqual(rename_stmt, expected_stmt)

    def test_drop_table_statement(self):
        parser = Parser()
        drop_queries = "DROP TABLE student_info"
        expected_stmt = DropObjectStatement(ObjectType.TABLE, "student_info", False)
        evadb_statement_list = parser.parse(drop_queries)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.DROP_OBJECT)
        drop_stmt = evadb_statement_list[0]
        self.assertEqual(drop_stmt, expected_stmt)

    def test_drop_udf_statement_str(self):
        drop_udf_query1 = """DROP UDF MyUDF;"""
        drop_udf_query2 = """DROP UDF IF EXISTS MyUDF;"""
        expected_stmt1 = DropObjectStatement(ObjectType.UDF, "MyUDF", False)
        expected_stmt2 = DropObjectStatement(ObjectType.UDF, "MyUDF", True)
        self.assertEqual(str(expected_stmt1), drop_udf_query1)
        self.assertEqual(str(expected_stmt2), drop_udf_query2)

    def test_single_statement_queries(self):
        parser = Parser()

        single_queries = []
        single_queries.append("SELECT CLASS FROM TAIPAI;")
        single_queries.append("SELECT CLASS FROM TAIPAI WHERE CLASS = 'VAN';")
        single_queries.append(
            "SELECT CLASS,REDNESS FROM TAIPAI \
            WHERE CLASS = 'VAN' AND REDNESS > 20.5;"
        )
        single_queries.append(
            "SELECT CLASS FROM TAIPAI \
            WHERE (CLASS = 'VAN' AND REDNESS < 300 ) OR REDNESS > 500;"
        )
        single_queries.append(
            "SELECT CLASS FROM TAIPAI \
            WHERE (CLASS = 'VAN' AND REDNESS < 300 ) OR REDNESS > 500;"
        )

        for query in single_queries:
            evadb_statement_list = parser.parse(query)

            self.assertIsInstance(evadb_statement_list, list)
            self.assertEqual(len(evadb_statement_list), 1)
            self.assertIsInstance(evadb_statement_list[0], AbstractStatement)

    def test_multiple_statement_queries(self):
        parser = Parser()

        multiple_queries = []
        multiple_queries.append(
            "SELECT CLASS FROM TAIPAI \
                WHERE (CLASS != 'VAN' AND REDNESS < 300)  OR REDNESS > 500; \
                SELECT REDNESS FROM TAIPAI \
                WHERE (CLASS = 'VAN' AND REDNESS = 300)"
        )

        for query in multiple_queries:
            evadb_statement_list = parser.parse(query)
            self.assertIsInstance(evadb_statement_list, list)
            self.assertEqual(len(evadb_statement_list), 2)
            self.assertIsInstance(evadb_statement_list[0], AbstractStatement)
            self.assertIsInstance(evadb_statement_list[1], AbstractStatement)

    def test_select_statement(self):
        parser = Parser()
        select_query = "SELECT CLASS, REDNESS FROM TAIPAI \
                WHERE (CLASS = 'VAN' AND REDNESS < 300 ) OR REDNESS > 500;"
        evadb_statement_list = parser.parse(select_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.SELECT)

        select_stmt = evadb_statement_list[0]

        # target List
        self.assertIsNotNone(select_stmt.target_list)
        self.assertEqual(len(select_stmt.target_list), 2)
        self.assertEqual(select_stmt.target_list[0].etype, ExpressionType.TUPLE_VALUE)
        self.assertEqual(select_stmt.target_list[1].etype, ExpressionType.TUPLE_VALUE)

        # from_table
        self.assertIsNotNone(select_stmt.from_table)
        self.assertIsInstance(select_stmt.from_table, TableRef)
        self.assertEqual(select_stmt.from_table.table.table_name, "TAIPAI")

        # where_clause
        self.assertIsNotNone(select_stmt.where_clause)
        # other tests should go in expression testing

    def test_select_with_empty_string_literal(self):
        parser = Parser()

        select_query = """SELECT '' FROM TAIPAI;"""

        evadb_statement_list = parser.parse(select_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.SELECT)

    def test_select_union_statement(self):
        parser = Parser()
        select_union_query = "SELECT CLASS, REDNESS FROM TAIPAI \
            UNION ALL SELECT CLASS, REDNESS FROM SHANGHAI;"
        evadb_statement_list = parser.parse(select_union_query)
        select_stmt = evadb_statement_list[0]
        self.assertIsNotNone(select_stmt.union_link)
        self.assertEqual(select_stmt.union_all, True)
        second_select_stmt = select_stmt.union_link
        self.assertIsNone(second_select_stmt.union_link)

    def test_select_statement_class(self):
        """Testing setting different clauses for Select
        Statement class
        Class: SelectStatement"""

        select_stmt_new = SelectStatement()
        parser = Parser()

        select_query_new = "SELECT CLASS, REDNESS FROM TAIPAI \
            WHERE (CLASS = 'VAN' AND REDNESS < 400 ) OR REDNESS > 700;"
        evadb_statement_list = parser.parse(select_query_new)
        select_stmt = evadb_statement_list[0]

        select_stmt_new.where_clause = select_stmt.where_clause
        select_stmt_new.target_list = select_stmt.target_list
        select_stmt_new.from_table = select_stmt.from_table

        self.assertEqual(select_stmt_new.where_clause, select_stmt.where_clause)
        self.assertEqual(select_stmt_new.target_list, select_stmt.target_list)
        self.assertEqual(select_stmt_new.from_table, select_stmt.from_table)
        self.assertEqual(str(select_stmt_new), str(select_stmt))

    def test_select_statement_groupby_class(self):
        """Testing sample frequency"""

        parser = Parser()

        select_query = "SELECT FIRST(id) FROM TAIPAI GROUP BY '8 frames';"

        evadb_statement_list = parser.parse(select_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.SELECT)

        select_stmt = evadb_statement_list[0]

        # target List
        self.assertIsNotNone(select_stmt.target_list)
        self.assertEqual(len(select_stmt.target_list), 1)
        self.assertEqual(
            select_stmt.target_list[0].etype, ExpressionType.AGGREGATION_FIRST
        )

        # from_table
        self.assertIsNotNone(select_stmt.from_table)
        self.assertIsInstance(select_stmt.from_table, TableRef)
        self.assertEqual(select_stmt.from_table.table.table_name, "TAIPAI")

        # sample_freq
        self.assertEqual(
            select_stmt.groupby_clause,
            ConstantValueExpression("8 frames", v_type=ColumnType.TEXT),
        )

    def test_select_statement_orderby_class(self):
        """Testing order by clause in select statement
        Class: SelectStatement"""

        parser = Parser()

        select_query = "SELECT CLASS, REDNESS FROM TAIPAI \
                    WHERE (CLASS = 'VAN' AND REDNESS < 400 ) OR REDNESS > 700 \
                    ORDER BY CLASS, REDNESS DESC;"
        # if orderby sort_type (ASC/DESC) not provided, should default to ASC

        evadb_statement_list = parser.parse(select_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.SELECT)

        select_stmt = evadb_statement_list[0]

        # target List
        self.assertIsNotNone(select_stmt.target_list)
        self.assertEqual(len(select_stmt.target_list), 2)
        self.assertEqual(select_stmt.target_list[0].etype, ExpressionType.TUPLE_VALUE)
        self.assertEqual(select_stmt.target_list[1].etype, ExpressionType.TUPLE_VALUE)

        # from_table
        self.assertIsNotNone(select_stmt.from_table)
        self.assertIsInstance(select_stmt.from_table, TableRef)
        self.assertEqual(select_stmt.from_table.table.table_name, "TAIPAI")

        # where_clause
        self.assertIsNotNone(select_stmt.where_clause)

        # orderby_clause
        self.assertIsNotNone(select_stmt.orderby_list)
        self.assertEqual(len(select_stmt.orderby_list), 2)
        self.assertEqual(select_stmt.orderby_list[0][0].name, "CLASS")
        self.assertEqual(select_stmt.orderby_list[0][1], ParserOrderBySortType.ASC)
        self.assertEqual(select_stmt.orderby_list[1][0].name, "REDNESS")
        self.assertEqual(select_stmt.orderby_list[1][1], ParserOrderBySortType.DESC)

    def test_select_statement_limit_class(self):
        """Testing limit clause in select statement
        Class: SelectStatement"""

        parser = Parser()

        select_query = "SELECT CLASS, REDNESS FROM TAIPAI \
                    WHERE (CLASS = 'VAN' AND REDNESS < 400 ) OR REDNESS > 700 \
                    ORDER BY CLASS, REDNESS DESC LIMIT 3;"

        evadb_statement_list = parser.parse(select_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.SELECT)

        select_stmt = evadb_statement_list[0]

        # target List
        self.assertIsNotNone(select_stmt.target_list)
        self.assertEqual(len(select_stmt.target_list), 2)
        self.assertEqual(select_stmt.target_list[0].etype, ExpressionType.TUPLE_VALUE)
        self.assertEqual(select_stmt.target_list[1].etype, ExpressionType.TUPLE_VALUE)

        # from_table
        self.assertIsNotNone(select_stmt.from_table)
        self.assertIsInstance(select_stmt.from_table, TableRef)
        self.assertEqual(select_stmt.from_table.table.table_name, "TAIPAI")

        # where_clause
        self.assertIsNotNone(select_stmt.where_clause)

        # orderby_clause
        self.assertIsNotNone(select_stmt.orderby_list)
        self.assertEqual(len(select_stmt.orderby_list), 2)
        self.assertEqual(select_stmt.orderby_list[0][0].name, "CLASS")
        self.assertEqual(select_stmt.orderby_list[0][1], ParserOrderBySortType.ASC)
        self.assertEqual(select_stmt.orderby_list[1][0].name, "REDNESS")
        self.assertEqual(select_stmt.orderby_list[1][1], ParserOrderBySortType.DESC)

        # limit_count
        self.assertIsNotNone(select_stmt.limit_count)
        self.assertEqual(select_stmt.limit_count, ConstantValueExpression(3))

    def test_select_statement_sample_class(self):
        """Testing sample frequency"""

        parser = Parser()

        select_query = "SELECT CLASS, REDNESS FROM TAIPAI SAMPLE 5;"

        evadb_statement_list = parser.parse(select_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.SELECT)

        select_stmt = evadb_statement_list[0]

        # target List
        self.assertIsNotNone(select_stmt.target_list)
        self.assertEqual(len(select_stmt.target_list), 2)
        self.assertEqual(select_stmt.target_list[0].etype, ExpressionType.TUPLE_VALUE)
        self.assertEqual(select_stmt.target_list[1].etype, ExpressionType.TUPLE_VALUE)

        # from_table
        self.assertIsNotNone(select_stmt.from_table)
        self.assertIsInstance(select_stmt.from_table, TableRef)
        self.assertEqual(select_stmt.from_table.table.table_name, "TAIPAI")

        # sample_freq
        self.assertEqual(select_stmt.from_table.sample_freq, ConstantValueExpression(5))

    def test_table_ref(self):
        """Testing table info in TableRef
        Class: TableInfo
        """
        table_info = TableInfo("TAIPAI", "Schema", "Database")
        table_ref_obj = TableRef(table_info)
        select_stmt_new = SelectStatement()
        select_stmt_new.from_table = table_ref_obj
        self.assertEqual(select_stmt_new.from_table.table.table_name, "TAIPAI")
        self.assertEqual(select_stmt_new.from_table.table.schema_name, "Schema")
        self.assertEqual(select_stmt_new.from_table.table.database_name, "Database")

    def test_insert_statement(self):
        parser = Parser()
        insert_query = """INSERT INTO MyVideo (Frame_ID, Frame_Path)
                                    VALUES    (1, '/mnt/frames/1.png');
                        """
        expected_stmt = InsertTableStatement(
            TableRef(TableInfo("MyVideo")),
            [
                TupleValueExpression("Frame_ID"),
                TupleValueExpression("Frame_Path"),
            ],
            [
                ConstantValueExpression(1),
                ConstantValueExpression("/mnt/frames/1.png", ColumnType.TEXT),
            ],
        )
        evadb_statement_list = parser.parse(insert_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.INSERT)

        insert_stmt = evadb_statement_list[0]
        self.assertEqual(insert_stmt, expected_stmt)

    def test_delete_statement(self):
        parser = Parser()
        delete_statement = """DELETE FROM Foo WHERE id > 5"""

        evadb_statement_list = parser.parse(delete_statement)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.DELETE)

        delete_stmt = evadb_statement_list[0]

        expected_stmt = DeleteTableStatement(
            TableRef(TableInfo("Foo")),
            ComparisonExpression(
                ExpressionType.COMPARE_GREATER,
                TupleValueExpression("id"),
                ConstantValueExpression(5),
            ),
        )

        self.assertEqual(delete_stmt, expected_stmt)

    def test_create_udf_statement(self):
        parser = Parser()
        create_udf_query = """CREATE UDF IF NOT EXISTS FastRCNN
                  INPUT  (Frame_Array NDARRAY UINT8(3, 256, 256))
                  OUTPUT (Labels NDARRAY STR(10), Bbox NDARRAY UINT8(10, 4))
                  TYPE  Classification
                  IMPL  'data/fastrcnn.py'
                  "KEY" "VALUE";
        """

        expected_cci = ColConstraintInfo()
        expected_cci.nullable = True
        expected_stmt = CreateUDFStatement(
            "FastRCNN",
            True,
            Path("data/fastrcnn.py"),
            [
                ColumnDefinition(
                    "Frame_Array",
                    ColumnType.NDARRAY,
                    NdArrayType.UINT8,
                    (3, 256, 256),
                    expected_cci,
                )
            ],
            [
                ColumnDefinition(
                    "Labels", ColumnType.NDARRAY, NdArrayType.STR, (10,), expected_cci
                ),
                ColumnDefinition(
                    "Bbox", ColumnType.NDARRAY, NdArrayType.UINT8, (10, 4), expected_cci
                ),
            ],
            "Classification",
            [("KEY", "VALUE")],
        )
        evadb_statement_list = parser.parse(create_udf_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.CREATE_UDF)
        self.assertEqual(str(evadb_statement_list[0]), str(expected_stmt))

        create_udf_stmt = evadb_statement_list[0]

        self.assertEqual(create_udf_stmt, expected_stmt)

    def test_load_video_data_statement(self):
        parser = Parser()
        load_data_query = """LOAD VIDEO 'data/video.mp4'
                             INTO MyVideo"""
        file_options = {}
        file_options["file_format"] = FileFormatType.VIDEO
        column_list = None
        expected_stmt = LoadDataStatement(
            TableInfo("MyVideo"),
            Path("data/video.mp4"),
            column_list,
            file_options,
        )
        evadb_statement_list = parser.parse(load_data_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.LOAD_DATA)

        load_data_stmt = evadb_statement_list[0]
        self.assertEqual(load_data_stmt, expected_stmt)

    def test_load_csv_data_statement(self):
        parser = Parser()
        load_data_query = """LOAD CSV 'data/meta.csv'
                             INTO
                             MyMeta (id, frame_id, video_id, label);"""
        file_options = {}
        file_options["file_format"] = FileFormatType.CSV
        expected_stmt = LoadDataStatement(
            TableInfo("MyMeta"),
            Path("data/meta.csv"),
            [
                TupleValueExpression("id"),
                TupleValueExpression("frame_id"),
                TupleValueExpression("video_id"),
                TupleValueExpression("label"),
            ],
            file_options,
        )
        evadb_statement_list = parser.parse(load_data_query)
        self.assertIsInstance(evadb_statement_list, list)
        self.assertEqual(len(evadb_statement_list), 1)
        self.assertEqual(evadb_statement_list[0].stmt_type, StatementType.LOAD_DATA)

        load_data_stmt = evadb_statement_list[0]
        self.assertEqual(load_data_stmt, expected_stmt)

    def test_nested_select_statement(self):
        parser = Parser()
        sub_query = """SELECT CLASS FROM TAIPAI WHERE CLASS = 'VAN'"""
        nested_query = """SELECT ID FROM ({}) AS T;""".format(sub_query)
        parsed_sub_query = parser.parse(sub_query)[0]
        actual_stmt = parser.parse(nested_query)[0]
        self.assertEqual(actual_stmt.stmt_type, StatementType.SELECT)
        self.assertEqual(actual_stmt.target_list[0].name, "ID")
        self.assertEqual(
            actual_stmt.from_table, TableRef(parsed_sub_query, alias=Alias("T"))
        )

        sub_query = """SELECT Yolo(frame).bbox FROM autonomous_vehicle_1
                              WHERE Yolo(frame).label = 'vehicle'"""
        nested_query = """SELECT Licence_plate(bbox) FROM
                            ({}) AS T
                          WHERE Is_suspicious(bbox) = 1 AND
                                Licence_plate(bbox) = '12345';
                      """.format(
            sub_query
        )
        query = """SELECT Licence_plate(bbox) FROM TAIPAI
                    WHERE Is_suspicious(bbox) = 1 AND
                        Licence_plate(bbox) = '12345';
                """
        query_stmt = parser.parse(query)[0]
        actual_stmt = parser.parse(nested_query)[0]
        sub_query_stmt = parser.parse(sub_query)[0]
        self.assertEqual(
            actual_stmt.from_table, TableRef(sub_query_stmt, alias=Alias("T"))
        )
        self.assertEqual(actual_stmt.where_clause, query_stmt.where_clause)
        self.assertEqual(actual_stmt.target_list, query_stmt.target_list)

    def test_should_return_false_for_unequal_expression(self):
        table = TableRef(TableInfo("MyVideo"))
        load_stmt = LoadDataStatement(
            table, Path("data/video.mp4"), FileFormatType.VIDEO
        )
        insert_stmt = InsertTableStatement(table)
        create_udf = CreateUDFStatement(
            "udf",
            False,
            Path("data/fastrcnn.py"),
            [
                ColumnDefinition(
                    "frame",
                    ColumnType.NDARRAY,
                    NdArrayType.UINT8,
                    (3, 256, 256),
                )
            ],
            [ColumnDefinition("labels", ColumnType.NDARRAY, NdArrayType.STR, (10))],
            "Classification",
        )
        select_stmt = SelectStatement()
        self.assertNotEqual(load_stmt, insert_stmt)
        self.assertNotEqual(insert_stmt, load_stmt)
        self.assertNotEqual(create_udf, insert_stmt)
        self.assertNotEqual(select_stmt, create_udf)

    def test_materialized_view(self):
        select_query = """SELECT id, Yolo(frame).labels FROM MyVideo
                        WHERE id<5; """
        query = "CREATE MATERIALIZED VIEW uadtrac_fastRCNN (id, labels) AS {}".format(
            select_query
        )
        parser = Parser()
        mat_view_stmt = parser.parse(query)
        select_stmt = parser.parse(select_query)
        expected_stmt = CreateMaterializedViewStatement(
            TableInfo("uadtrac_fastRCNN"),
            [
                ColumnDefinition("id", None, None, None),
                ColumnDefinition("labels", None, None, None),
            ],
            False,
            select_stmt[0],
        )
        self.assertEqual(mat_view_stmt[0], expected_stmt)

    def test_join(self):
        select_query = """SELECT table1.a FROM table1 JOIN table2
                    ON table1.a = table2.a; """
        parser = Parser()
        select_stmt = parser.parse(select_query)[0]
        table1_col_a = TupleValueExpression("a", "table1")
        table2_col_a = TupleValueExpression("a", "table2")
        select_list = [table1_col_a]
        from_table = TableRef(
            JoinNode(
                TableRef(TableInfo("table1")),
                TableRef(TableInfo("table2")),
                predicate=ComparisonExpression(
                    ExpressionType.COMPARE_EQUAL, table1_col_a, table2_col_a
                ),
                join_type=JoinType.INNER_JOIN,
            )
        )
        expected_stmt = SelectStatement(select_list, from_table)

        self.assertEqual(select_stmt, expected_stmt)

    def test_join_with_where(self):
        select_query = """SELECT table1.a FROM table1 JOIN table2
            ON table1.a = table2.a WHERE table1.a <= 5"""
        parser = Parser()
        select_stmt = parser.parse(select_query)[0]
        table1_col_a = TupleValueExpression("a", "table1")
        table2_col_a = TupleValueExpression("a", "table2")
        select_list = [table1_col_a]
        from_table = TableRef(
            JoinNode(
                TableRef(TableInfo("table1")),
                TableRef(TableInfo("table2")),
                predicate=ComparisonExpression(
                    ExpressionType.COMPARE_EQUAL, table1_col_a, table2_col_a
                ),
                join_type=JoinType.INNER_JOIN,
            )
        )
        where_clause = ComparisonExpression(
            ExpressionType.COMPARE_LEQ,
            table1_col_a,
            ConstantValueExpression(5),
        )
        expected_stmt = SelectStatement(select_list, from_table, where_clause)
        self.assertEqual(select_stmt, expected_stmt)

    def test_multiple_join_with_multiple_ON(self):
        select_query = """SELECT table1.a FROM table1 JOIN table2
            ON table1.a = table2.a JOIN table3
            ON table3.a = table1.a WHERE table1.a <= 5"""
        parser = Parser()
        select_stmt = parser.parse(select_query)[0]
        table1_col_a = TupleValueExpression("a", "table1")
        table2_col_a = TupleValueExpression("a", "table2")
        table3_col_a = TupleValueExpression("a", "table3")
        select_list = [table1_col_a]
        child_join = TableRef(
            JoinNode(
                TableRef(TableInfo("table1")),
                TableRef(TableInfo("table2")),
                predicate=ComparisonExpression(
                    ExpressionType.COMPARE_EQUAL, table1_col_a, table2_col_a
                ),
                join_type=JoinType.INNER_JOIN,
            )
        )

        from_table = TableRef(
            JoinNode(
                child_join,
                TableRef(TableInfo("table3")),
                predicate=ComparisonExpression(
                    ExpressionType.COMPARE_EQUAL, table3_col_a, table1_col_a
                ),
                join_type=JoinType.INNER_JOIN,
            )
        )
        where_clause = ComparisonExpression(
            ExpressionType.COMPARE_LEQ,
            table1_col_a,
            ConstantValueExpression(5),
        )
        expected_stmt = SelectStatement(select_list, from_table, where_clause)
        self.assertEqual(select_stmt, expected_stmt)

    def test_lateral_join(self):
        select_query = """SELECT frame FROM MyVideo JOIN LATERAL
                            ObjectDet(frame) AS OD;"""
        parser = Parser()
        select_stmt = parser.parse(select_query)[0]
        tuple_frame = TupleValueExpression("frame")
        func_expr = FunctionExpression(
            func=None, name="ObjectDet", children=[tuple_frame]
        )
        from_table = TableRef(
            JoinNode(
                TableRef(TableInfo("MyVideo")),
                TableRef(TableValuedExpression(func_expr), alias=Alias("OD")),
                join_type=JoinType.LATERAL_JOIN,
            )
        )
        expected_stmt = SelectStatement([tuple_frame], from_table)
        self.assertEqual(select_stmt, expected_stmt)

    def test_class_equality(self):
        table_info = TableInfo("MyVideo")
        table_ref = TableRef(TableInfo("MyVideo"))
        tuple_frame = TupleValueExpression("frame")
        func_expr = FunctionExpression(
            func=None, name="ObjectDet", children=[tuple_frame]
        )
        join_node = JoinNode(
            TableRef(TableInfo("MyVideo")),
            TableRef(TableValuedExpression(func_expr), alias=Alias("OD")),
            join_type=JoinType.LATERAL_JOIN,
        )
        self.assertNotEqual(table_info, table_ref)
        self.assertNotEqual(tuple_frame, table_ref)
        self.assertNotEqual(join_node, table_ref)
        self.assertNotEqual(table_ref, table_info)

    def test_lark(self):
        query = """CREATE UDF FaceDetector
                  INPUT  (frame NDARRAY UINT8(3, ANYDIM, ANYDIM))
                  OUTPUT (bboxes NDARRAY FLOAT32(ANYDIM, 4),
                          scores NDARRAY FLOAT32(ANYDIM))
                  TYPE  FaceDetection
                  IMPL  'evadb/udfs/face_detector.py';
                  """
        parser = Parser()
        parser.parse(query)


import uuid

import six

from tooz import coordination

coordinator = coordination.get_coordinator('kazoo://localhost', b'host-1')
coordinator.start()

# Create a group
group = six.binary_type(six.text_type(uuid.uuid4()).encode('ascii'))
request = coordinator.create_group(group)
request.get()

# Join a group
request = coordinator.join_group(group)
request.get()

coordinator.stop()


import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ExpSineSquared, ConstantKernel, WhiteKernel

class GPReg:
    def __init__(self):
        self.kernel = 2**2*RBF(0.5, (1e-1, 1e0)) + WhiteKernel(0.01,(0.09999,0.010001))
        #self.kernel = 4.**2*ExpSineSquared(1.0, 2*np.pi, (1e-5, 1e5), (2*3.11,2*3.15))

    def test(self,UX,UY,X):
        UX = np.array(UX)
        UY = np.array(UY)
        X = np.array(X)
        y_pred = np.zeros((X.shape[0],X.shape[1],UY.shape[2]))
        sigma = np.zeros((X.shape[0],X.shape[1],UY.shape[2], UY.shape[2]))

        for i in range(X.shape[0]):
            x = X[i,:,:]
            ux = UX[i,:,:]
            uy = UY[i,:,:]

            # independent GP for each output dim
            for j in range(UY.shape[2]):
                gp = GaussianProcessRegressor(kernel=self.kernel, n_restarts_optimizer=9)
                if ux.shape[0] > 0:
                    gp.fit(ux,uy[:,j])

                y,s = gp.predict(x, return_std=True)
                #y = np.reshape(y,[-1, UY.shape[2]])
                y_pred[i,:,j], sigma[i,:,j,j] = (y,s)

        return y_pred, sigma


import posixpath

from django.core.urlresolvers import reverse

class Finder(object):
    template_base = ""
    columns = []

    def __init__(self, auth=None):
        self.auth = auth
        self.columns_by_name = dict((c.name, c) for c in self.columns)
        self.parent_columns = dict(
            zip([c.name for c in self.columns[1:]], self.columns[:-1])
            )
        self.child_columns = dict(
            zip([c.name for c in self.columns[:-1]], self.columns[1:])
            )
        self.columns_by_model = dict((c.model, c) for c in self.columns)

    def column_template(self, column_name):
        """
        Returns the template to use for rendering a list of objects in a given
        column.

        """
        col = self._get_column_by_name(column_name)
        return posixpath.join(self.template_base, col.template_name)

    def goto_url(self, obj):
        """
        Given an object, return its "Goto" url, or None.

        """
        try:
            col = self.columns_by_model[obj.__class__]
        except KeyError:
            return None

        if col.goto is not None:
            return reverse(col.goto) + "?%s=%s" % (col.attr, obj.id)

        return None

    def child_column_for_obj(self, obj):
        """
        Given an object, return the column name of its child column, or None if
        there is none.

        """
        try:
            col = self.columns_by_model[obj.__class__]
            return self.child_columns[col.name].name
        except KeyError:
            return None

    def child_query_url(self, obj):
        """
        Given an object, return the URL for an ajax query to fetch child
        objects.

        """
        child_col = self.child_column_for_obj(obj)
        if child_col is not None:
            return "?finder=1&col=%s&id=%s" % (child_col, obj.id)
        return None

    def objects(self, column_name, parent_id=None):
        """
        Given a column name, return the list of objects.

        If a parent_id is given and there is a parent column, filter the list
        by that parent id.

        """
        col = self._get_column_by_name(column_name)
        ret = col.objects(auth=self.auth)
        if parent_id is not None:
            try:
                parent = self.parent_columns[col.name]
            except KeyError:
                raise ValueError(
                    "Column %r has no parent; cannot use parent_id."
                    % column_name)
            else:
                return ret.filter(**{parent.attr: parent_id})
        return ret

    def _get_column_by_name(self, column_name):
        try:
            return self.columns_by_name[column_name]
        except KeyError:
            raise ValueError("Column %r does not exist." % column_name)

class Column(object):
    def __init__(self, name, template_name, list_model, attr,
                 goto=None, sort=("name", "asc"), **filters):
        self.name = name
        self.template_name = template_name
        self.list_model = list_model
        self.model = self.list_model.entryclass
        self.attr = attr
        self.filters = filters
        self.sort = sort
        self.goto = goto

    def objects(self, auth):
        return self.list_model.ours(auth=auth).filter(
            **self.filters).sort(*self.sort)


# ===========================================================================
#
# file     : gf_max7219.py
# part of  : godafoss micropython library
# url      : https://www.github.com/wovo/godafoss
# author   : Wouter van Ooijen (wouter@voti.nl) 2023
# license  : MIT license, see license variable in the __init__.py
#
# This file is part of the Godafoss perhiperal interface library.
#
# This file contains the max7219 LED matrix driver class.
#
# ===========================================================================

from micropython import const
import framebuf

from godafoss.gf_xy import *
from godafoss.gf_pins import *
from godafoss.gf_make_pins import *
from godafoss.gf_canvas import *

# ===========================================================================

class max7219( canvas ):
    """
    max7219 LED matrix driver

    The max7219 is a LED matrix driver that is commonly used to drive
    either an 8x8 LED matrix, or up to 8 7-segment LED displays.
    Max7219 chips can be chained to drive a larger LED matrix.

    The
    """

    # =======================================================================

    class _commands:
        NOOP          = const( 0x00 )
        DIGIT_ZERO    = const( 0x01 ) # + digit (y) number 0..7
        DECODE_MODE   = const( 0x09 ) # data: 0bxxxx_xxxx mode of each digit
        INTENSITY     = const( 0x0A ) # data: 0bxxxx duty cycle
        SCAN_LIMIT    = const( 0x0B ) # data: 0bxxx 0-7 for 1-8 displays
        SHUT_DOWN     = const( 0x0C ) # data: 0b1 = normal, 0b0 = shutdown
        DISPLAY_TEST  = const( 0x0F ) # data: 0b0 = normal, 0b1 = test

    # =======================================================================

    def __init__(
        self,
        n,
        spi,
        chip_select : [ int, pin_out, pin_in_out, pin_oc ],
        background = False,
        brightness = 0,
        enable = True
    ):
        self._n = n
        self._spi = spi
        self._chip_select = make_pin_out( chip_select )

        canvas.__init__(
            self,
            size = xy( 8 * self._n, 8 ),
            is_color = False,
            background = background
        )

        self.buffer = bytearray(( self.size.y // 8 ) * self.size.x )
        self.framebuf = framebuf.FrameBuffer(
            self.buffer, self.size.x, self.size.y, framebuf.MONO_HLSB
        )

        self.write_command(  self._commands.DISPLAY_TEST, 0 ),
        self.write_command(  self._commands.SCAN_LIMIT, 7 ),
        self.write_command(  self._commands.DECODE_MODE, 0 ),

        self.brightness( brightness )
        self.enable( enable )

    # =======================================================================

    def enable( self, v ):
        self.write_command(
            self._commands.SHUT_DOWN,
            0x01 if v else 0x00
        )

    # =======================================================================

    def brightness( self, v ):
        self.write_command(
            self._commands.INTENSITY,
            clamp( v, 0, 7 )
        )

    # =======================================================================

    def _write_pixel_implementation(
        self,
        location: xy,
        ink: bool | None = True
    ) -> None:
        self.framebuf.pixel(
            location.x,
            location.y,
            ink
        )

    # =======================================================================

    def write_command(
        self,
        command,
        data
    ) -> None:
        self._chip_select.write( 0 )
        for m in range( self._n ):
            self._spi.write( bytearray( [ command, data ] ) )
        self._chip_select.write( 1 )

    # =======================================================================

    def _flush_implementation( self ) -> None:

        for y in range( 8 ):
            self._chip_select.write( 0 )
            for x8 in range( self._n ):
                self._spi.write( bytearray( [
                    self._commands.DIGIT_ZERO + y ,
                    self.buffer[ ( y * self._n ) + x8 ]
                ] ) )
            self._chip_select.write( 1 )

    # =======================================================================

# ===========================================================================


import uuid
import time
import logging
import mongoengine as me

from mist.api.helpers import is_email_valid
from mist.api.logs.methods import log_event
from mist.api.users.models import User
from mist.api.machines.models import Machine

log = logging.getLogger(__name__)

ACTIONS = {}  # This is a map of action types to action classes.

def _populate_actions():
    """Populate ACTIONS variable."""
    for key, value in globals().iteritems():
        if key.endswith('Action') and key != 'Action':
            if issubclass(value, BaseAlertAction):
                if value.atype not in (None, 'no_data', ):  # Exclude these.
                    ACTIONS[value.atype] = value

class BaseAlertAction(me.EmbeddedDocument):
    """The base class for all alert actions.

    This class serves as very basic, common interface amongst subclasses that
    define alert actions. Every subclass of the `BaseAlertAction` MUST define
    at least its own `run` method, which holds all the logic of the action's
    execution, and a descriptive `atype`.

    A rule may define a list of actions to be executed once it's triggered. The
    actions will be sequentially executed in the order they have been provided.

    """

    meta = {'allow_inheritance': True}

    atype = None

    id = me.StringField(required=True, default=lambda: uuid.uuid4().hex)

    def update(self, fail_on_error=True, **kwargs):
        for key, value in kwargs.iteritems():
            if key not in self._fields:
                if not fail_on_error:
                    continue
                raise me.ValidationError('Field "%s" does not exist on %s',
                                         key, type(self))
            setattr(self, key, value)

    def run(self):
        """Execute self.

        The body of the action to be executed. Subclasses MUST override this.

        """
        raise NotImplementedError()

    def as_dict(self):
        return {'type': self.atype}

    def __str__(self):
        return '%s %s' % (self.__class__.__name__, self.id)

class NotificationAction(BaseAlertAction):
    """An action that notifies the users, once a rule has been triggered."""

    atype = 'notification'

    users = me.ListField(me.StringField(), default=lambda: [])
    teams = me.ListField(me.StringField(), default=lambda: [])
    emails = me.ListField(me.StringField(), default=lambda: [])

    def run(self, machine, value, triggered, timestamp, incident_id, action='',
            notification_level=0):
        # FIXME Imported here due to circular dependency issues.
        from mist.api.notifications.methods import send_alert_email
        # TODO Shouldn't be specific to machines.
        assert isinstance(machine, Machine)
        assert machine.owner == self._instance.owner
        emails = set(self.emails)
        user_ids = set(self.users)
        if not (self.users or self.teams):
            emails |= set(self._instance.owner.get_emails())
            emails |= set(self._instance.owner.alerts_email)
        if user_ids:
            user_ids &= set([m.id for m in self._instance.owner.members])
        for user in User.objects(id__in=user_ids):
            emails.add(user.email)
        for team_id in self.teams:
            try:
                team = self._instance.owner.teams.get(id=team_id)
                emails |= set([member.email for member in team.members])
            except me.DoesNotExist:
                continue
        send_alert_email(machine.owner, self._instance.id, value,
                         triggered, timestamp, incident_id, emails,
                         cloud_id=machine.cloud.id,
                         machine_id=machine.machine_id, action=action)

    def clean(self):
        """Perform e-mail address validation."""
        for email in self.emails:
            if not is_email_valid(email):
                raise me.ValidationError('Invalid e-mail address: %s' % email)

    def as_dict(self):
        return {'type': self.atype, 'emails': self.emails,
                'users': self.users, 'teams': self.teams}

class NoDataAction(NotificationAction):
    """An action triggered in case of a NoData alert."""

    atype = 'no_data'

    def run(self, machine, value, triggered, timestamp, incident_id, **kwargs):
        if timestamp + 60 * 60 * 24 < time.time():
            # FIXME Imported here due to circular dependency issues.
            from mist.api.monitoring.methods import disable_monitoring
            # If NoData alerts are being triggered for over 24h, disable
            # monitoring and log the action to close any open incidents.
            disable_monitoring(machine.owner, machine.cloud.id,
                               machine.machine_id, no_ssh=True)
            log_event(
                machine.owner.id, 'incident', 'disable_monitoring',
                cloud_id=machine.cloud.id, machine_id=machine.machine_id,
                incident_id=incident_id
            )
            action = 'Disable Monitoring'
            notification_level = 0
        else:
            action = 'Alert'
            notification_level = kwargs.get('notification_level', 0)
        super(NoDataAction, self).run(machine, value, triggered, timestamp,
                                      incident_id, action, notification_level)

class CommandAction(BaseAlertAction):
    """Execute a remote command."""

    # TODO: Deprecate in favor of a ScriptAction?

    atype = 'command'

    command = me.StringField(required=True)

    def run(self, machine, *args, **kwargs):
        # FIXME Imported here due to circular dependency issues.
        from mist.api.methods import ssh_command
        assert isinstance(machine, Machine)
        assert machine.owner == self._instance.owner
        return ssh_command(machine.owner, machine.cloud.id, machine.machine_id,
                           machine.hostname, self.command)

    def as_dict(self):
        return {'type': self.atype, 'command': self.command}

class MachineAction(BaseAlertAction):
    """Perform a machine action."""

    atype = 'machine_action'

    action = me.StringField(required=True, choices=('reboot', 'destroy'))

    def run(self, machine, *args, **kwargs):
        assert isinstance(machine, Machine)
        assert machine.owner == self._instance.owner
        getattr(machine.ctl, self.action)()
        if self.action == 'destroy':  # If destroy, disable monitoring, too.
            # FIXME Imported here due to circular dependency issues.
            from mist.api.monitoring.methods import disable_monitoring
            # TODO Move this into machine.ctl.destroy method and
            # deprecate mist.api.machines.methods:destroy_machine.
            # Could also be implemented as new method inside the
            # MachineController.
            disable_monitoring(machine.owner, machine.cloud.id,
                               machine.machine_id, no_ssh=True)

    def as_dict(self):
        return {'type': self.atype, 'action': self.action}

_populate_actions()


# coding: utf-8
# 匿名函数

### -> Syntax：
# lambda x:x+1  x 函数形参   x+1 函数返回值

### -> Demo_1
# def calc1(x):
#     return x + 1
# a = calc1(10)
# print(a)

# # 等于

# calc2 = lambda x:x+1  # <------ 定义函数并赋值给 a， 此时 func: calc1== func: calc2
# a = calc2(10)
# print(a)

### -> Demo_2
name = "alex"

ch_name = lambda x:x + "_sb"
print(ch_name(name))


#!/usr/bin/env python
# ----------------------------------------------------------------------------
# Copyright (c) 2016-2018, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------

import biom
import warnings
import seaborn as sns
import pandas as pd
import numpy as np
import q2templates
import pkg_resources

from shutil import copyfile
from os.path import join
from q2_types.tree import NewickFormat
from itolapi import Itol

TEMPLATES = pkg_resources.resource_filename('q2_qemistree', 'assets')

def values_to_colors(coloring_category: str, color_palette: str):
    '''This function generates a color map (dict) for unique values in a
    user-specified feature metadata column.'''
    color_map = {}
    annotations = coloring_category.unique()
    colors = sns.color_palette(color_palette,
                               n_colors=len(annotations)).as_hex()
    # give a heads up to the user
    if len(set(colors)) < len(annotations):
        warnings.warn("The mapping between colors and annotations"
                      " is not unique, some colors have been repeated",
                      UserWarning)
    for i, value in enumerate(annotations):
        color_map[value] = colors[i]
    return color_map

def format_colors(feature_metadata, category, color_palette):
    colors = []

    color_map = values_to_colors(feature_metadata[category], color_palette)

    colors.append('TREE_COLORS')
    colors.append('SEPARATOR TAB')
    colors.append('DATA')

    for idx in feature_metadata.index:
        color = color_map[feature_metadata.loc[idx, category]]

        if feature_metadata.loc[idx, 'structure_source'] == 'MS2':
            style, width = 'normal', 6
        else:
            style, width = 'dashed', 4

        colors.append('%s\tclade\t%s\t%s\t%s' % (idx, color, style, width))

    return '\n'.join(colors)

def format_labels(feature_metadata, category, ms2_label, parent_mz):
    labels = []

    missing_values = {'unclassified', 'unexpected server response',
                      'SMILE parse error', np.nan}

    labels.append('LABELS')
    labels.append('SEPARATOR TAB')
    labels.append('DATA')

    if ms2_label:
        for idx in feature_metadata.index:
            ms2_compound = feature_metadata.loc[idx, 'ms2_library_match']
            if pd.notna(ms2_compound) and not ms2_compound.isspace():
                label = ms2_compound
            else:
                label = feature_metadata.loc[idx, category]

            if parent_mz and label in missing_values:
                label = feature_metadata.loc[idx, parent_mz]

            labels.append('%s\t%s' % (idx, label))
    else:
        for idx in feature_metadata.index:
            label = feature_metadata.loc[idx, category]

            if parent_mz and label in missing_values:
                label = feature_metadata.loc[idx, parent_mz]

            labels.append('%s\t%s' % (idx, label))

    return '\n'.join(labels)

def plot(output_dir: str, table: biom.Table, tree: NewickFormat,
         feature_metadata: pd.DataFrame, category: str,
         color_palette: str = 'Dark2', ms2_label: bool = False,
         parent_mz: str = None) -> None:

    if category not in feature_metadata.columns:
        raise ValueError('Could not find %s in the feature data, the available'
                         ' columns are: %s.' %
                         (category,
                          ', '.join(feature_metadata.columns.astype())))

    color_fp = join(output_dir, 'colors.tsv')
    with open(color_fp, 'w') as fh:
        fh.write(format_colors(feature_metadata, category, color_palette))

    label_fp = join(output_dir, 'labels.tsv')
    with open(label_fp, 'w') as fh:
        fh.write(format_labels(feature_metadata, category, ms2_label,
                               parent_mz))

    # itol won't accept a file unless it has a .tree or .txt extension
    target = join(output_dir, 'qemistree.tree')
    copyfile(str(tree), target)

    # upload the tree, labels and tip coloring files
    itol_uploader = Itol()
    itol_uploader.add_file(target)
    itol_uploader.add_file(label_fp)
    itol_uploader.add_file(color_fp)

    status = itol_uploader.upload()

    if not status:
        raise ValueError(itol_uploader.comm.upload_output)

    url = itol_uploader.get_webpage()
    print(url)

    index_path = join(TEMPLATES, 'index.html')
    q2templates.render(index_path, output_dir, context={'url': url})


#
# @lc app=leetcode id=7 lang=python3
#
# [7] Reverse Integer
#
# https://leetcode.com/problems/reverse-integer/description/
#
# algorithms
# Easy (25.62%)
# Likes:    2791
# Dislikes: 4370
# Total Accepted:    931.7K
# Total Submissions: 3.6M
# Testcase Example:  '123'
#
# Given a 32-bit signed integer, reverse digits of an integer.
#
# Example 1:
#
#
# Input: 123
# Output: 321
#
#
# Example 2:
#
#
# Input: -123
# Output: -321
#
#
# Example 3:
#
#
# Input: 120
# Output: 21
#
#
# Note:
# Assume we are dealing with an environment which could only store integers
# within the 32-bit signed integer range: [−2^31,  2^31 − 1]. For the purpose
# of this problem, assume that your function returns 0 when the reversed
# integer overflows.
#
#

# @lc code=start 方法一：str
# class Solution:
#     def reverse(self, x: int) -> int:
#         sol = lambda x: x if x in range (-2**31-1,2**31) else 0
#         if (x in range (-2**31-1,2**31)):
#             if (x < 0):
#                 return sol(-int(str(x)[:0:-1]))
#             else:
#                 return sol(int(str(x)[::-1]))
#         else:
#             return 0

# @lc code=start 方法二：math
class Solution:
    def reverse(self, x: int) -> int:
        sign = -1 if x < 0 else 1
        inrange = lambda x: x if x in range (-2**31-1,2**31) else 0
        if inrange(x)==0:
            return 0
        else:
            ans = 0
            x = abs(x)
            while(x):
                ans = ans*10+x%10
                x=x//10
            return inrange(ans)*sign

# @lc code=end


#!/usr/bin/python
# coding: utf-8

# Copyright (c) 2016 Mountainstorm
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import unicode_literals, print_function

# 0x4D, 0x69, 0x63, 0x72, 0x6F, 0x73, 0x6F, 0x66, // "Microsof"
# 0x74, 0x20, 0x43, 0x2F, 0x43, 0x2B, 0x2B, 0x20, // "t C/C++ "
# 0x4D, 0x53, 0x46, 0x20, 0x37, 0x2E, 0x30, 0x30, // "MSF 7.00"
# 0x0D, 0x0A, 0x1A, 0x44, 0x53, 0x00, 0x00, 0x00  // "^^^DS^^^"

class PdbFileHeader(object):
    def __init__(self, reader, bits):
        bits.min_capacity(56)
        reader.seek(0)
        bits.fill_buffer(reader, 52)

        self.magic = bits.read_bytes(32)            #   0..31
        self.page_size = bits.read_int32()          #  32..35
        self.free_pagemap = bits.read_int32()       #  36..39
        self.pages_used = bits.read_int32()         #  40..43
        self.directory_size = bits.read_int32()     #  44..47
        self.zero = bits.read_int32()               #  48..51

        directory_pages = (
            ((((self.directory_size + self.page_size - 1) / self.page_size) *
            4) + self.page_size - 1) / self.page_size
        )
        bits.fill_buffer(reader, directory_pages * 4)
        self.directory_root = bits.read_int32(directory_pages)


from flask import Flask,render_template,jsonify
from flask_cors import CORS
from traitement_question import Question
from geocoding import Geocoding
from wiki import Wiki

app = Flask(__name__)
CORS(app)

@app.route('/')
def index():
  return render_template('index.html')

@app.route('/question/<user_question>',methods=['get'])
def question(user_question):
  answer = Question.traitement_question(user_question)
  geo = Geocoding.data_geocoding(answer)
  ville = (geo[3])
  wiki = Wiki.wiki_story(ville)

  return jsonify({"status":"OK" if geo else "error","address":geo[0] if geo else None,"lon":geo[2] if geo else None,"lat":geo[1] if geo else None,"wiki":wiki,"user_question_sw":answer})

if __name__ == '__main__':
    app.run(debug=False)


class DataController:
    bot = None
    users = {}
    emoji = {
        "0": "0⃣",
        "1": "1⃣",
        "2": "2⃣",
        "3": "3⃣",
        "4": "4⃣",
        "5": "5⃣",
        "6": "6⃣",
        "7": "7⃣",
        "8": "8⃣",
        "9": "9⃣"
    }

    def __init__(self, bot):
        self.bot = bot
        bot.add_cog(self)

        for y in bot.get_all_emojis():
            self.emoji[y.name] = y

        for y in bot.get_all_members():
            self.users[str(y)] = self.User(y)

    async def on_message(self, message):
        if str(message.author) in self.users:
            member = self.users[str(message.author)]
            self.users[str(message.author)].last_message = member.message
            self.users[str(message.author)].message = str(message.content)
        else:
            raise Exception("Unknown user " + str(message.author))

    async def on_member_join(self, member):
        self.users[str(member)] = self.User(member)
        print("User joined")

    class User:
        def __init__(self, user):
            self.message = ""
            self.last_message = ""
            self.user = user


#Jariel Long
#CS5590
#Lab 1

#1
tup = ( 'John', ('Physics', 80)) , ('Daniel', ('Science', 90)), \
      ('John', ('Science', 95)), ('Mark',('Maths', 100)), \
      ('Daniel', ('History', 75)), ('Mark', ('Social', 95))

z = {}
for x,y in tup:
    if x not in z:
        z[x] = []
        z[x].append(y)
    else:
        z[x].append(y)
        z[x].sort()

for k,v in z.items():
    print(k, v)

#2
###### make permutations

x = "pwwkew"
tset = set()
tempstr = ""
for y in x:
    '''add elements to set'''
    tset.add(y)
for y in tset:
    '''combine set into string'''
    tempstr = tempstr+y

from itertools import permutations
#find permutations
l = list(permutations(tempstr, len(tempstr)))
for x in l:
    #print permutations
    print(x , len(tempstr))

#3   AirLine
class flight:
    """A class representing a flight"""
    def __init__(self, plane, time, destination):
        self.plane = plane
        self.time = time
        self.destination = destination

class person(flight):
    #Inhereitance
    """A class representing a person"""
    def __init__(self, plane, time, destination, name, seat, ssn):
        flight.__init__(self, plane, time, destination)
        self.name = name
        self.seat = seat
        #ssn is obviously private
        self.__ssn = ssn

class employee(person):
    #inhereiteance
    """A class representing a employee"""
    def __init__(self, plane, time, destination, name, seat, ssn, title):
        #super call
        person.__init__(self, plane,time, destination, name, seat , ssn)
        self.title = title

#Instances
fl = flight(1, 22, "America")
pe = person(2, 20, "America", "Bob", "23A", "850-97-5698")
em = employee(3, 11, "Paris", "Jen", "FrontCoach", "687-98-7412", "Pilot")

#4
import pandas as pd
import numpy as np
from sklearn.metrics import r2_score
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn import linear_model
from sklearn import svm
from sklearn.naive_bayes import GaussianNB
from sklearn.neighbors import KNeighborsClassifier

#had to change diabetes positive to 1
#and diabetes negative to 0
#original and changed csv are in folder
train = pd.read_csv("diabetes2_csv.csv")

numeric_features = train.select_dtypes(include = [np.number])
#calculate correlations
corr = numeric_features.corr()
print(corr['class'].sort_values(ascending=False)[:4], '\n') #top 4 correlated
print(corr['class'].sort_values(ascending=False)[-8:]) # bottom 8 correlated
#Fill null values
data = train.select_dtypes(include=[np.number]).interpolate().dropna()
X_train = data.drop(['pres','skin','class'], axis=1) #more accurate
#X_train = data.drop(['class'],axis=1) #a little less accurate
y_train = train['class']
#Splits dataset into train and test
X_train, X_test, y_train, y_test = train_test_split(X_train, y_train, test_size = 0.3, random_state=0)
lrl = linear_model.LinearRegression()

model = lrl.fit(X_train, y_train)
pred = model.predict(X_test)
print('rmse', mean_squared_error(y_test, pred))
print('r2_Score', r2_score(y_test, pred))

#Before the least correlated features are dropped the rmse value is close to 0, however the average value of both rmse and r2 are closer to 0 when the least correlated features are kept
#indicating that the predicated values are very close and in line with the observed values
#r2 score claims that the dependent variables are explained by the collective independent variables

#5
#had to change >50k to 0
#had to change <=50k to 1
#Had to change an extensive amount of data such as race for example to numerical values

train = pd.read_csv("adult.csv")
numeric_features = train.select_dtypes(include = [np.number])
#calculate correlations
corr = numeric_features.corr()
print(corr['inc'].sort_values(ascending=False)[:5], '\n') #top 4 correlated
print(corr['inc'].sort_values(ascending=False)[-12:]) # bottom 8 correlated

#Fill null values
data = train.select_dtypes(include=[np.number]).interpolate().dropna()
#drop low correlated values
X_train = data.drop(["ed-num","age","hpw", "cg", "sex", "cl", "fnlwgt", "race", "inc"],axis=1) #more accurate
y_train = train['inc']
#Splits dataset into train and test
X_train, X_test, y_train, y_test = train_test_split(X_train, y_train, test_size = 0.4, random_state=0)
#SVM
svc = svm.SVC()
svc.fit(X_train, y_train)
Y_pred = svc.predict(X_test)
acc_svc = round(svc.score(X_train, y_train)*100,2)
print("svm accuracy is: ", acc_svc)
#NB
gnb = GaussianNB()
gnb.fit(X_train,y_train)
y_pred = gnb.predict(X_test)
acc_gnb = round(gnb.score(X_train, y_train)* 100,2)
print("gnb accuracy is: ", acc_gnb)
#KNN
knn = KNeighborsClassifier(n_neighbors =3)
knn.fit(X_train, y_train)
y_pred = knn.predict(X_test)
acc_knn = round(knn.score(X_train, y_train)*100,2)
print("KNN Accuracy is: ", acc_knn)

#SVM has the best accuracy followed by KNN and GNB in last

#6
from sklearn.decomposition import PCA
from sklearn.preprocessing import LabelEncoder, StandardScaler
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn import preprocessing
from sklearn import metrics

train = pd.read_csv('wiki4HE.csv')
x = train
#fill null with mean
x = x.fillna(x.mean())

#standardrize the data
scaler = preprocessing.StandardScaler()
scaler.fit(x)
X_scaled_array = scaler.transform(x)
X_scaled = pd.DataFrame(X_scaled_array, columns = x.columns)

km = KMeans(n_clusters=5, random_state=0)

km.fit(X_scaled)

y_cluster_kmeans = km.predict(X_scaled)

#Silhouette Score
y_cluster_kmeans = km.predict(x)
score = metrics.silhouette_score(x, y_cluster_kmeans)
print(score)
#Score is about .3683

#elbow method to know the number of clusters
wcss = []
for i in range(1,11):
    kmeans = KMeans(n_clusters=i, init='k-means++', max_iter=300, n_init=10, random_state=0)
    kmeans.fit(x)
    wcss.append(kmeans.inertia_)

plt.plot(range(1, 11), wcss)
plt.title('the elbow method')
plt.xlabel('Number of Clusters')
plt.ylabel('Wcss')
plt.show()
#Elbow is about 4 or 5

pca = PCA(2)
x_pca = pca.fit_transform(X_scaled_array)
df2 = pd.DataFrame(data=x_pca)
finaldf = pd.concat([df2],axis=1)

print(finaldf)

km.fit(finaldf)

# predict the cluster for each data point
y_cluster_kmeans = km.predict(finaldf)
score = metrics.silhouette_score(finaldf, y_cluster_kmeans)
print(score)


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

numberOfPaymentsM=[] #number of the months
monthlyInstallmentM=[] # money to pay per month [$usd]
interestM=[] # interest to pay per month [$usd]
aInterestM=[] #Cumulative interest to pay per month [$usd]
principalM=[] # Principal matrix [$usd]
aPrincipalM=[] # Cumulative Principal matrix [$usd]
balanceM=[] #balance per month [$USD]

def PlotLoan(numberOfPaymentsM,aInterestM,aPrincipalM,balanceM):
    plt.figure(1)
    plt.plot(numberOfPaymentsM,balanceM,'r',numberOfPaymentsM,aInterestM,'b',numberOfPaymentsM,aPrincipalM,'g')
   # plt.show()

def Mycsv2(periods):
    with open('Loan.csv','w',newline='') as f:
        fieldnames=['Number of payments','Monthly installment','Interest','Principal','Balance']
        #fieldnames=['Period','Monthly energy production','Panel degradation','Energy generation','Minimum Fees','Cashflow','Net cashflow']
        thewriter=csv.DictWriter(f ,fieldnames=fieldnames)
        thewriter.writeheader()

        for i in range (periods+1):
            number=numberOfPaymentsM[i]
            installment=monthlyInstallmentM[i]
            interest=interestM[i]
            principal=principalM[i]
            balance=balanceM[i]
            thewriter.writerow({'Number of payments':number,'Monthly installment':installment,'Interest':interest,'Principal':principal,'Balance':balance})
            #thewriter.writerow({'Period':period, 'Monthly energy production':pro,'Panel degradation': deg,'Energy generation': gen,'Minimum Fees':fee,'Cashflow':cas,'Net cashflow':net})

def loan(rate,periods,quantity):

    installment = np.pmt(rate,periods,-quantity) # monthly installment
    numberOfPaymentsM.append(0)
    monthlyInstallmentM.append(0)
    interestM.append(0)
    principalM.append(0)
    aInterestM.append(0)
    aPrincipalM.append(0)
    balanceM.append(quantity)

    for i in range(1,periods+1):
        numberOfPaymentsM.append(i)
        monthlyInstallmentM.append(installment)
        interest=balanceM[i-1]*rate
        interestM.append(interest)
        principal=monthlyInstallmentM[i]-interestM[i]
        principalM.append(principal)
        balance=balanceM[i-1]-principalM[i]
        balanceM.append(balance)
        acumi=sum(interestM)
        aInterestM.append(acumi)
        aprin=sum(principalM)
        aPrincipalM.append(aprin)

    PlotLoan(numberOfPaymentsM,aInterestM,aPrincipalM,balanceM)
    Mycsv2(periods)

    return (monthlyInstallmentM)


#%%

import torch
import matplotlib.pyplot as plt
import torch.nn as nn
import numpy as np
from scipy.stats import norm
import util

# device

print(f'Is CUDA available?: {torch.cuda.is_available()}')
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# device = 'cpu'

# hyperparameters

torch.manual_seed(123)
trun_N = 4
dim_xi = 10
layer_sizes_ubar = [2,32,32,32,1]
layer_sizes_U = [2,64,64,64,trun_N]
layer_sizes_A = [1,4,4,4,1]
layer_sizes_Y = [1+dim_xi,64,64,64,trun_N]
activation = torch.tanh
step_size = 0.001
epochs = 101
a = 0.5
b = 0.3
sigma_g = 1
lc = 0.4
method = 'BO'

xl, xu = -1.0, 1.0
tl, tu = 0.0, 1.0
# nx, nt, nxi = 51, 50, 1000
# nvalx, nvalt, nvalxi = 51, 50, 1000
nx, nt, nxi = 50, 30, 200
nvalx, nvalt, nvalxi = 50, 30, 200

# data

def ubar_t0(x):
    return -torch.sin(np.pi*x)
def u_t0(x):
    return torch.sin(np.pi*(torch.arange(trun_N).to(device)+1).view(1,-1)*(x+1)/2)
def Y_t0(xi):
    return xi[:,:trun_N]
def a_t0(t):
    return torch.zeros_like(t).repeat(1,trun_N)

data = {}
data['x'] = torch.linspace(xl,xu,nx)
data['xb'] = torch.Tensor([xl,xu])
data['t'] = (tu-tl)*torch.rand(nt)+tl
data['t0'] = torch.Tensor([tl])
data['xi'] = torch.randn([nxi,dim_xi])
data['wxi'] = 1/nxi*torch.ones(nxi)
data['wx'] = (xu-xl)/nx*torch.ones(nx)
data['wt'] = (tu-tl)/nt*torch.ones(nt)
data['wt0'] = torch.Tensor([1.])
data['xEu'] = torch.Tensor([-0.5,0.0,0.5])
data['tEu'] = torch.Tensor([0.1,0.9])

data_val = {}
data_val['x'] = torch.linspace(xl,xu,nvalx)
data_val['xb'] = torch.Tensor([xl,xu])
data_val['t'] = torch.linspace(tl,tu,nvalt)
data_val['t0'] = torch.Tensor([tl])
data_val['xi'] = torch.randn([nvalxi,dim_xi])
data_val['wxi'] = 1/nvalxi*torch.ones(nvalxi)
data_val['wx'] = (xu-xl)/nvalx*torch.ones(nvalx)
data_val['wt'] = (tu-tl)/nvalt*torch.ones(nvalt)
data_val['wt0'] = torch.Tensor([1.])
data_val['xEu'] = torch.Tensor([-0.5,0.0,0.5])
data_val['tEu'] = torch.Tensor([0.1,0.9])

for d in data:
    data[d] = data[d].to(device)
for d in data_val:
    data_val[d] = data_val[d].to(device)

# KLD of GP

KL_z, KL_w = np.polynomial.legendre.leggauss(dim_xi+20)
covmat = sigma_g**2 * np.exp( -(KL_z.reshape(-1,1) - KL_z.reshape(1,-1))**2 / lc**2 )
covmat_data = sigma_g**2 * np.exp( -(data['x'].cpu().numpy().reshape(-1,1) - KL_z.reshape(1,-1))**2 / lc**2 )
covmat_val = sigma_g**2 * np.exp( -(data_val['x'].cpu().numpy().reshape(-1,1) - KL_z.reshape(1,-1))**2 / lc**2 )

sqrtW = np.diag(np.sqrt(KL_w))
WcovmatW = np.dot(np.dot(sqrtW,covmat),sqrtW)
lamk, uk = np.linalg.eig(WcovmatW)
lamk = np.real(lamk[:dim_xi])
uk = np.real(uk[:,:dim_xi])
fk_data = 1/np.sqrt(lamk).reshape(1,-1)*np.dot(covmat_data,uk*np.sqrt(KL_w).reshape(-1,1))
fk_val = 1/np.sqrt(lamk).reshape(1,-1)*np.dot(covmat_val,uk*np.sqrt(KL_w).reshape(-1,1))
fk_data = torch.Tensor(fk_data).to(device)
fk_val = torch.Tensor(fk_val).to(device)
f_data = (1 - data['x'].view(1,-1)**2) * (1 + torch.mm(data['xi'],fk_data.T))
f_val = (1 - data_val['x'].view(1,-1)**2) * (1 + torch.mm(data_val['xi'],fk_val.T))

# exact for u

xxi = torch.cartesian_prod(data_val['x'],data_val['xi'][:,0])[:,0:1]
xxi = torch.cat((xxi,data_val['xi'].repeat(nvalx,1)),1)
Nht = 100
ht = (tu-tl)/(nvalt-1)/Nht
hx = (xu-xl)/(nvalx-1)
exact_u = ubar_t0(xxi[:,0:1]).view(nvalx,1,nvalxi)
D2 = torch.zeros(nvalx,nvalx).to(device)
for i in range(nvalx):
    for j in range(nvalx):
        if i==0:
            D2[i,i] = 1.
        elif i==nvalx-1:
            D2[i,i] = 1.
        else:
            if i==j:
                D2[i,j] = -2/hx**2
            elif abs(i-j)==1:
                D2[i,j] = 1/hx**2

for i in range((nvalt-1)*Nht):
    if i<2:
        tmp = a*torch.matmul(D2,exact_u[:,i,:]) + b*exact_u[:,i,:]**2 + f_val.T
        tmp = exact_u[:,i,:] + ht*tmp
        exact_u = torch.cat((exact_u,tmp.view(nvalx,1,nvalxi)),1)
    else:
        tmp1 = a*torch.matmul(D2,exact_u[:,i,:]) + b*exact_u[:,i,:]**2 + f_val.T
        tmp2 = a*torch.matmul(D2,exact_u[:,i-1,:]) + b*exact_u[:,i-1,:]**2 + f_val.T
        tmp3 = a*torch.matmul(D2,exact_u[:,i-2,:]) + b*exact_u[:,i-2,:]**2 + f_val.T
        tmp = exact_u[:,i,:] + ht/12 * (23*tmp1 - 16*tmp2 + 5*tmp3)
        exact_u = torch.cat((exact_u,tmp.view(nvalx,1,nvalxi)),1)

exact_u = exact_u[:,::Nht,:]

xEutEu = torch.cartesian_prod(data_val['xEu'],data_val['tEu']).cpu().numpy()

inv_x = (xEutEu[:,0] - xl)/(xu - xl)*(nvalx-1)
inv_t = (xEutEu[:,1] - tl)/(tu - tl)*(nvalt-1)
inv_x = tuple(map(int,inv_x))
inv_t = tuple(map(int,inv_t))
data_u = exact_u[inv_x,inv_t,:]
data_u = data_u * data_val['wxi'].view(1,nvalxi)
data_u = data_u.sum(1)

data['Eu'] = data_u
data_val['Eu'] = data_u

# model

class Net4(nn.Module):

    def __init__(self, layer_sizes, activation=torch.tanh):
        super(Net4, self).__init__()
        self.layer_sizes = layer_sizes
        self.layer_list = []
        self.activation = activation

        self.l1 = nn.Linear(layer_sizes[0], layer_sizes[1])
        self.l2 = nn.Linear(layer_sizes[1], layer_sizes[2])
        self.l3 = nn.Linear(layer_sizes[2], layer_sizes[3])
        self.l4 = nn.Linear(layer_sizes[3], layer_sizes[4])

    def forward(self, x):
        x = self.l1(x)
        x = self.activation(x)
        x = self.l2(x)
        x = self.activation(x)
        x = self.l3(x)
        x = self.activation(x)
        x = self.l4(x)
        return x

net_ubar = Net4(layer_sizes_ubar, activation).to(device)
net_U = Net4(layer_sizes_U, activation).to(device)
nets = [net_ubar,net_U]
for _ in range(trun_N):
    nets.append(Net4(layer_sizes_A, activation).to(device))
net_Y = Net4(layer_sizes_Y, activation).to(device)
nets.append(net_Y)
n_params_single = 2

def model_loss(data, fmodel, params_unflattened, tau_likes, gradients, params_single=None):

    xt = torch.cartesian_prod(data['x'],data['t'])
    wxt = torch.cartesian_prod(data['wx'],data['wt'])
    t = data['t'].view(-1,1)
    wt = data['wt'].view(-1,1)
    x = data['x'].view(-1,1)
    wx = data['wx'].view(-1,1)
    xi = data['xi']
    wxi = data['wxi'].view(-1,1)
    t0 = data['t0'].view(-1,1)
    wt0 = data['wt0'].view(-1,1)
    xt0 = torch.cartesian_prod(data['x'],data['t0'])
    wxt0 = torch.cartesian_prod(data['wx'],data['wt0'])
    xbt = torch.cartesian_prod(data['xb'],data['t'])
    xEutEu = torch.cartesian_prod(data['xEu'],data['tEu'])
    Eu = data['Eu'].view(-1,1)

    Nx = x.shape[0]
    Nt = t.shape[0]
    Nxi = xi.shape[0]

    xtxi = torch.cartesian_prod(data['x'],data['t'],data['xi'][:,0])[:,0:2]
    xtxi = torch.cat((xtxi,data['xi'].repeat(Nx*Nt,1)),1)
    wxtxi = torch.cartesian_prod(data['wx'],data['wt'],data['wxi'])

    txi = torch.cartesian_prod(data['t'],data['xi'][:,0])[:,0:1]
    txi = torch.cat((txi,data['xi'].repeat(Nt,1)),1)
    wtxi = torch.cartesian_prod(data['wt'],data['wxi'])

    t0xi = torch.cat((t0.repeat(Nxi,1),xi),1)
    wt0xi = torch.cartesian_prod(data['wt0'],data['wxi'])

    # MSE_w

    xtxi = xtxi.detach().requires_grad_()
    ubar = fmodel[0](xtxi[:,(0,1)], params=params_unflattened[0])
    U = fmodel[1](xtxi[:,(0,1)], params=params_unflattened[1])
    A = fmodel[2](xtxi[:,1:2], params=params_unflattened[2])
    for i in range(trun_N-1):
        A = torch.cat((A,fmodel[i+3](xtxi[:,1:2], params=params_unflattened[i+3])),1)
    Y = fmodel[trun_N+2](xtxi[:,1:], params=params_unflattened[trun_N+2])

    u = ubar + (A*U*Y).sum(1).unsqueeze(1)
    Du = gradients(u,xtxi)[0]
    u_x, u_t = Du[:,0:1], Du[:,1:2]
    u_xx = gradients(u_x,xtxi)[0][:,0:1]
    Nxu = params_single[0]*u_xx + params_single[1]*u**2 + ( f_data.T.view(Nx,1,Nxi) + torch.zeros(1,Nt,1).to(device) ).reshape(-1,1)

    tmp = (u_t - Nxu)*wxtxi[:,2:3].prod(1).unsqueeze(1)
    tmp = tmp.view(Nx,Nt,Nxi).sum(2)**2
    tmp = tmp.view(Nx*Nt,1)*wxt.prod(1).unsqueeze(1)
    MSE_w = tmp.sum()

    tmp = (u_t - Nxu)*U*wxtxi[:,0:1]
    tmp = tmp.view(Nx,Nt,Nxi,trun_N).sum(0)**2
    tmp = tmp.view(Nt*Nxi,trun_N)*wtxi.prod(1).unsqueeze(1)
    MSE_w += tmp.sum()

    tmp = (u_t - Nxu)*Y*wxtxi[:,2:3].prod(1).unsqueeze(1)
    tmp = tmp.view(Nx,Nt,Nxi,trun_N).sum(2)**2
    tmp = tmp.view(Nx*Nt,trun_N)*wxt.prod(1).unsqueeze(1)
    MSE_w += tmp.sum()

    # MSE_0

    tmp = u_t - Nxu
    tmp = tmp**2 * wxtxi.prod(1).unsqueeze(1)
    MSE_0 = tmp.sum()

    output = [u,U,A,Y,ubar]

    # MSE_inv
    ubar = fmodel[0](xEutEu, params=params_unflattened[0])
    tmp = ubar - Eu
    tmp = tmp**2
    MSE_inv = tmp.sum()

    # MSE_IC

    ubar = fmodel[0](xt0, params=params_unflattened[0])
    U = fmodel[1](xt0, params=params_unflattened[1])
    A = fmodel[2](t0, params=params_unflattened[2])
    for i in range(trun_N-1):
        A = torch.cat((A,fmodel[i+3](t0, params=params_unflattened[i+3])),1)
    Y = fmodel[trun_N+2](t0xi, params=params_unflattened[trun_N+2])

    tmp = ubar - ubar_t0(xt0[:,0:1])
    tmp = tmp**2 * wxt0.prod(1).unsqueeze(1)
    MSE_IC = tmp.sum()

    tmp = U - u_t0(xt0[:,0:1])
    tmp = tmp**2 * wxt0.prod(1).unsqueeze(1)
    MSE_IC += tmp.sum()/trun_N

    tmp = A - a_t0(t0)
    tmp = tmp**2 * wt0.prod(1).unsqueeze(1)
    MSE_IC += tmp.sum()/trun_N

    tmp = Y - Y_t0(t0xi[:,1:])
    tmp = tmp**2 * wt0xi.prod(1).unsqueeze(1)
    MSE_IC += tmp.sum()/trun_N

    # MSE_BC

    ubar = fmodel[0](xbt, params=params_unflattened[0])
    U = fmodel[1](xbt, params=params_unflattened[1])
    A = fmodel[2](t, params=params_unflattened[2])
    for i in range(trun_N-1):
        A = torch.cat((A,fmodel[i+3](t, params=params_unflattened[i+3])),1)
    Y = fmodel[trun_N+2](txi, params=params_unflattened[trun_N+2])

    tmp = ubar.view(-1,Nt)
    tmp = tmp**2 * wt.view(1,Nt)
    MSE_BC = tmp.sum()

    tmp = Y.view(Nt,Nxi,trun_N) * torch.sqrt(wxi.prod(1).view(1,Nxi,1))
    CYY = torch.matmul(torch.transpose(tmp,1,2),tmp) # shape: (Nt,trun_N,trun_N)
    tmp = U.view(-1,Nt*trun_N)
    tmp = tmp.T.view(Nt,trun_N,-1)
    tmp = tmp * A.view(Nt,trun_N,1)
    tmp = torch.matmul(CYY,tmp)
    tmp = tmp.view(Nt,-1)**2 * wt
    MSE_BC += tmp.sum()/trun_N

    # MSE_DO

    if method == 'DO':

        xt = xt.detach().requires_grad_()
        txi = txi.detach().requires_grad_()
        U = fmodel[1](xt, params=params_unflattened[1])
        Y = fmodel[trun_N+2](txi, params=params_unflattened[trun_N+2])

        tmp = Y*wtxi[:,1:].prod(1).unsqueeze(1)
        tmp = tmp.view(Nt,Nxi,trun_N).sum(1)**2
        tmp = tmp.view(Nt,trun_N) * wt
        MSE_DO = tmp.sum()/trun_N

        U_t = torch.zeros(Nx*Nt,trun_N).to(device)
        for i in range(trun_N):
            U_t[:,i:i+1] = gradients(U[:,i:i+1],xt)[0][:,1:2]
        U = U.view(Nx,Nt,trun_N)
        U_t = U_t.view(Nx,Nt,trun_N)
        U = torch.transpose(U,0,1) * wx.view(1,Nx,1) # shape: (Nt,Nx,trun_N)
        U_t = torch.transpose(U_t,0,1) # shape: (Nt,Nx,trun_N)
        tmp = torch.matmul(torch.transpose(U_t,1,2),U) # shape: (Nt,trun_N,trun_N)
        tmp = tmp**2 * wt.view(Nt,1,1)
        MSE_DO += tmp.sum()/trun_N**2

        Y_t = torch.zeros(Nt*Nxi,trun_N).to(device)
        for i in range(trun_N):
            Y_t[:,i:i+1] = gradients(Y[:,i:i+1],txi)[0][:,0:1]
        Y = Y.view(Nt,Nxi,trun_N)
        Y_t = Y_t.view(Nt,Nxi,trun_N)
        tmp = Y*Y_t*wxi.prod(1).view(1,Nxi,1)
        tmp = tmp.sum(1)**2 * wt.view(Nt,1)
        MSE_DO += tmp.sum()/trun_N

    elif method == 'BO':

        xt = xt.detach().requires_grad_()
        txi = txi.detach().requires_grad_()
        U = fmodel[1](xt, params=params_unflattened[1])
        Y = fmodel[trun_N+2](txi, params=params_unflattened[trun_N+2])

        tmp = Y*wtxi[:,1:]
        tmp = tmp.view(Nt,Nxi,trun_N).sum(1)**2
        tmp = tmp.view(Nt,trun_N) * wt
        MSE_BO = tmp.sum()/trun_N

        U_t = torch.zeros(Nx*Nt,trun_N).to(device)
        for i in range(trun_N):
            U_t[:,i:i+1] = gradients(U[:,i:i+1],xt)[0][:,1:2]
        U = U.view(Nx,Nt,trun_N)
        U_t = U_t.view(Nx,Nt,trun_N)
        U = torch.transpose(U,0,1) * wx.view(1,Nx,1) # shape: (Nt,Nx,trun_N)
        U_t = torch.transpose(U_t,0,1) # shape: (Nt,Nx,trun_N)
        tmp = torch.matmul(torch.transpose(U_t,1,2),U) # shape: (Nt,trun_N,trun_N)
        tmp = tmp + torch.transpose(tmp,1,2)
        tmp = tmp**2 * wt.view(Nt,1,1)
        MSE_BO += tmp.sum()/trun_N**2

        Y_t = torch.zeros(Nt*Nxi,trun_N).to(device)
        for i in range(trun_N):
            Y_t[:,i:i+1] = gradients(Y[:,i:i+1],txi)[0][:,0:1]
        Y = Y.view(Nt,Nxi,trun_N) * wxi.prod(1).view(1,Nxi,1)
        Y_t = Y_t.view(Nt,Nxi,trun_N)
        tmp = torch.matmul(torch.transpose(Y,1,2),Y_t) # shape: (Nt,trun_N,trun_N)
        tmp = tmp + torch.transpose(tmp,1,2)
        tmp = tmp**2 * wt.view(Nt,1,1)
        MSE_BO += tmp.sum()/trun_N

    if method == 'DO':
        ll = MSE_w + 100 * (MSE_IC + MSE_BC + MSE_DO + MSE_inv) + 0.1*MSE_0
    elif method == 'BO':
        ll = MSE_w + 100 * (MSE_IC + MSE_BC + MSE_BO + MSE_inv) + 0.1*MSE_0
    ll = -ll

    if torch.cuda.is_available():
        del xtxi, wxtxi, xt, wxt, txi, wtxi, t, wt, wx, wxi, t0, xt0, wxt0, t0xi, wt0xi, xbt, xEutEu, Eu, ubar, U, A, Y, u, Du, u_x, u_xx, u_t, Nxu, tmp, CYY, U_t, Y_t, MSE_w, MSE_IC, MSE_BC, MSE_0, MSE_inv
        if method == 'DO':
            del MSE_DO
        elif method == 'BO':
            del MSE_BO
        torch.cuda.empty_cache()

    return ll, output

# regression

params_init_val = None

params_hmc = util.sample_model_bpinns(nets, data, model_loss=model_loss, step_size=step_size, device=device, n_params_single=n_params_single, pde=True, pinns=True, epochs=epochs, params_init_val = params_init_val)

pred_list, log_prob_list = util.predict_model_bpinns(nets, params_hmc, data_val, model_loss=model_loss, n_params_single=n_params_single, pde = True)

print('\nExpected validation log probability: {:.3f}'.format(torch.stack(log_prob_list).mean()))

params_single = torch.stack(params_hmc)[:,:n_params_single].cpu().numpy()

print('Inferred a,b: {}'.format(params_single.mean(0)))

pred_list_u = pred_list[0].cpu().numpy()
pred_list_U = pred_list[1].cpu().numpy()
pred_list_A = pred_list[2].cpu().numpy()
pred_list_Y = pred_list[3].cpu().numpy()
pred_list_ubar = pred_list[4].cpu().numpy()

exact_u = exact_u.cpu().numpy()

# plot

wxi = data_val['wxi'].cpu().numpy()
x = data_val['x'].cpu().numpy()
t = data_val['t'].cpu().numpy()

pred_mean_u = pred_list_u.mean(0).reshape(nvalx,nvalt,nvalxi)
pred_mean_u = pred_mean_u * wxi.reshape(1,1,nvalxi)
pred_mean_u = pred_mean_u.sum(2)

pred_std_u = pred_list_u.mean(0).reshape(nvalx,nvalt,nvalxi)
pred_std_u = pred_std_u**2 * wxi.reshape(1,1,nvalxi)
pred_std_u = pred_std_u.sum(2) - pred_mean_u**2
pred_std_u = pred_std_u**0.5

exact_mean_u = exact_u * wxi.reshape(1,1,nvalxi)
exact_mean_u = exact_mean_u.sum(2)

exact_std_u = exact_u**2 * wxi.reshape(1,1,nvalxi)
exact_std_u = exact_std_u.sum(2) - exact_mean_u**2
exact_std_u = exact_std_u**0.5

pred_A = pred_list_A.mean(0).reshape(nvalx,nvalt,nvalxi,trun_N)[0,:,0,:]

pred_ubar = pred_list_ubar.mean(0).reshape(nvalx,nvalt,nvalxi)[:,:,0]

# graph in paper

att1 = int(1*(nvalt-1)/10)
att2 = int(10*(nvalt-1)/10)

plt.figure()
plt.plot(x,pred_mean_u[:,att1],marker='s',markersize=5, label='T=0.1, NN-'+method)
plt.plot(x,pred_mean_u[:,att2],marker='^',markersize=5, label='T=1.0, NN-'+method)
plt.plot(x,exact_mean_u[:,att1],'r--', label='T=0.1, Exact')
plt.plot(x,exact_mean_u[:,att2],'r-.', label='T=1.0, Exact')
plt.legend(fontsize=10)

plt.figure()
plt.plot(x,pred_std_u[:,att1]**2,marker='s',markersize=5, label='T=0.1, NN-'+method)
plt.plot(x,pred_std_u[:,att2]**2,marker='^',markersize=5, label='T=1.0, NN-'+method)
plt.plot(x,exact_std_u[:,att1]**2,'r--', label='T=0.1, Exact')
plt.plot(x,exact_std_u[:,att2]**2,'r-.', label='T=1.0, Exact')
plt.legend(fontsize=10)

plt.figure()
plt.plot(t,abs(pred_A),marker='s',markersize=5, label='NN-'+method)

plt.figure()
plt.plot(x,pred_mean_u[:,att2],'b-', label='T=1, mean, NN-'+method)
plt.fill_between(x,pred_mean_u[:,att2] - pred_std_u[:,att2], pred_mean_u[:,att2] + pred_std_u[:,att2],'b', label='T=1, std, NN-'+method, alpha=0.2)
plt.plot(x,exact_mean_u[:,att2],'r-', label='T=1, mean, reference')
plt.fill_between(x,exact_mean_u[:,att2] - exact_std_u[:,att2], exact_mean_u[:,att2] + exact_std_u[:,att2],'r', label='T=1, std, reference', alpha=0.2)
plt.legend(fontsize=10)

# %%


# bfs_kbacon.py
"""Volume 2A: Breadth-First Search (Kevin Bacon).
<Juan>
<345>
<Today>
"""

from collections import deque
import networkx as nx
from matplotlib import pyplot as plt

# Problems 1-4: Implement the following class
class Graph(object):
    """A graph object, stored as an adjacency dictionary. Each node in the
    graph is a key in the dictionary. The value of each key is a list of the
    corresponding node's neighbors.

    Attributes:
        dictionary: the adjacency list of the graph.
    """

    def __init__(self, adjacency):
        """Store the adjacency dictionary as a class attribute."""
        self.dictionary = adjacency

    # Problem 1
    def __str__(self):
        """String representation: a sorted view of the adjacency dictionary.

        Example:
            >>> test = {'A':['B'], 'B':['A', 'C',], 'C':['B']}
            >>> print(Graph(test))
            A: B
            B: A; C
            C: B
        """
        total = list()
        for i in xrange(len(self.dictionary)):
            letters = str()
            self.dictionary[self.dictionary.keys()[i]].sort()
            for j in xrange(len(self.dictionary[self.dictionary.keys()[i]])):
                if j == 0:
                    letters += str(self.dictionary[str(self.dictionary.keys()[i])][j])
                else:
                    letters += "; " + str(self.dictionary[str(self.dictionary.keys()[i])][j])
            total.append(str(str(self.dictionary.keys()[i]) + ": " + letters))
        total.sort()
        return '\n'.join(total)

    # Problem 2
    def traverse(self, start):
        """Begin at 'start' and perform a breadth-first search until all
        nodes in the graph have been visited. Return a list of values,
        in the order that they were visited.

        Inputs:
            start: the node to start the search at.

        Returns:
            the list of visited nodes (in order of visitation).

        Raises:
            ValueError: if 'start' is not in the adjacency dictionary.

        Example:
            >>> test = {'A':['B'], 'B':['A', 'C',], 'C':['B']}
            >>> Graph(test).traverse('B')
            ['B', 'A', 'C']
        """
        if start not in self.dictionary:
            raise ValueError("That is not in the adjacency")
        current = start

        marked = set()
        visited = list()
        visit_queue = deque()

        adjacency_dictionary = self.dictionary
        marked.add(current)

        for n in xrange(len(self.dictionary)-1):
            for neighbor in adjacency_dictionary[current]:
                if neighbor not in marked:
                    visit_queue.append(neighbor)
                    marked.add(neighbor)
            # print current
            visited.append(current)
            current = visit_queue.popleft()
            n -= 1
        visited.append(current)

        return (visited)

    # Problem 3 (Optional)
    def DFS(self, start):
        """Begin at 'start' and perform a depth-first search until all
        nodes in the graph have been visited. Return a list of values,
        in the order that they were visited. If 'start' is not in the
        adjacency dictionary, raise a ValueError.

        Inputs:
            start: the node to start the search at.

        Returns:
            the list of visited nodes (in order of visitation)
        """
        raise NotImplementedError("Problem 3 Incomplete")

    # Problem 4
    def shortest_path(self, start, target):
        """Begin at the node containing 'start' and perform a breadth-first
        search until the node containing 'target' is found. Return a list
        containg the shortest path from 'start' to 'target'. If either of
        the inputs are not in the adjacency graph, raise a ValueError.

        Inputs:
            start: the node to start the search at.
            target: the node to search for.

        Returns:
            A list of nodes along the shortest path from start to target,
                including the endpoints.

        Example:
            >>> test = {'A':['B', 'F'], 'B':['A', 'C'], 'C':['B', 'D'],
            ...         'D':['C', 'E'], 'E':['D', 'F'], 'F':['A', 'E', 'G'],
            ...         'G':['A', 'F']}
            >>> Graph(test).shortest_path('A', 'G')
            ['A', 'F', 'G']
        """
        if start not in self.dictionary:
            raise ValueError(start, "is not in the adjacency")

        if target not in self.dictionary:
            raise ValueError(target," is not in the adjacency")

        current = start

        new_dict = dict()

        marked = set()
        visited = list()
        visit_queue = deque()

        adjacency_dictionary = self.dictionary
        marked.add(current)
        while current != target:
            for neighbor in adjacency_dictionary[current]:
                if neighbor not in marked:
                    visit_queue.append(neighbor)
                    new_dict[neighbor] = current
                    marked.add(neighbor)
            # print current
            visited.append(current)
            current = visit_queue.popleft()

        visited.append(current)

        #path
        last = target
        path = []
        while last != start:
            path.append(last)
            last = new_dict[last]
        path.append(start)
        path.reverse()
        return path

# Problem 5: Write the following function
def convert_to_networkx(dictionary):
    """Convert 'dictionary' to a networkX object and return it."""
    nx_graph = nx.Graph()
    for x in dictionary.keys():
        for y in dictionary[x]:
            nx_graph.add_edge(x, y)

    return nx_graph

# Helper function for problem 6
def parse(filename="movieData.txt"):
    """Generate an adjacency dictionary where each key is
    a movie and each value is a list of actors in the movie.
    """

    # open the file, read it in, and split the text by '\n'
    with open(filename, 'r') as movieFile:
        moviesList = movieFile.read().split('\n')
    graph = dict()

    # for each movie in the file,
    for movie in moviesList:
        # get movie name and list of actors
        names = movie.split('/')
        title = names[0]
        graph[title] = []
        # add the actors to the dictionary
        for actor in names[1:]:
            graph[title].append(actor)

    return graph

# Problems 6-8: Implement the following class
class BaconSolver(object):
    """Class for solving the Kevin Bacon problem."""

    # Problem 6
    def __init__(self, filename="movieData.txt"):
        """Initialize the networkX graph and with data from the specified
        file. Store the graph as a class attribute. Also store the collection
        of actors in the file as an attribute.
        """
        actores = set()
        self.filename = filename
        dictionary = parse(filename)
        for x in dictionary.keys():
            for y in dictionary[x]:
                actores.add(y)
        self.actors = actores
        self.graph = convert_to_networkx(dictionary)

    # Problem 6
    def path_to_bacon(self, start, target="Bacon, Kevin"):
        """Find the shortest path from 'start' to 'target'."""
        if start not in self.actors:
            raise ValueError(start, "sucks and is not a friend of Kevin")
        if target not in self.actors:
            raise ValueError("Coder: you should have stayed with Kevin")

        return nx.shortest_path(self.graph, start, target)

    # Problem 7
    def bacon_number(self, start, target="Bacon, Kevin"):
        """Return the Bacon number of 'start'."""
        number = (len(nx.shortest_path(self.graph, start, target))-1)/2.
        return number

    # Problem 7
    def average_bacon(self, target="Bacon, Kevin"):
        """Calculate the average Bacon number in the data set.
        Note that actors are not guaranteed to be connected to the target.

        Inputs:
            target (str): the node to search the graph for
        """
        number = 0
        not_friends = 0
        for actor in self.actors:
            try :
                number += self.bacon_number(actor, target)
            except :
                not_friends += 1

        average = number/float(len(self.actors)-not_friends)
        return average, not_friends

# =========================== END OF FILE =============================== #

"""
if __name__ == "__main__":
    #prob1
    test = {'F':['G', 'E', 'A'],'B':['A', 'Z', 'C'], 'C':['T', 'R', 'B', 'D'],'D':['C', 'E'], 'A':['B', 'F'],  'E':['D', 'F'],
     'G':['A', 'F']}
    graph = Graph(test)
    #print graph

    #prob2
    test = {'A':['B', 'C', 'D', 'E'], 'B':['A', 'C'],
'C':['B', 'A', 'D'], 'D':['A', 'C'], 'E':['A']}
    #print Graph(test).traverse('B')

    #prob3
    Graph(test).shortest_path("D", "D")

    #prob4
    convert_to_networkx(test)

    #prob5
    movie_graph = BaconSolver("movieData.txt")
    movie_graph.path_to_bacon('Stark, Peter')
    movie_graph.bacon_number('Stark, Peter')
    movie_graph.average_bacon()

    my_dictionary = {'A':['C', 'B'], 'C':['A'], 'B':['A']}
    graph = Graph(my_dictionary)
    print(graph)
    print
    test =  {'A':['B', 'C', 'D', 'E'], 'B':['A', 'C'],
'C':['B', 'A', 'D'], 'D':['A', 'C'], 'E':['A']}
    print Graph(test).traverse('D')
    print

    test = {'A':['B', 'F'], 'B':['A', 'C'], 'C':['B', 'D'],'D':['C', 'E'], 'E':['D', 'F'], 'F':['A', 'E', 'G'],'G':['A', 'F']}
    print Graph(test).shortest_path('A','G')

    bacon = BaconSolver()

    movie_graph = BaconSolver("movieData.txt")
    print(movie_graph.path_to_bacon("Lee, Reggie"))
    print(movie_graph.bacon_number("Lee, Reggie"))
    print(movie_graph.average_bacon())

"""