from __future__ import print_function as _
from __future__ import division as _
from __future__ import absolute_import as _
import unittest
import numpy as np
import copy
import turicreate as tc
from turicreate.toolkits._main import ToolkitError
from turicreate.toolkits._private_utils import _validate_lists
from pandas.util.testing import assert_frame_equal
from . import util
import array
import string
import random
import sys
from functools import partial

from .test_distances import euclidean, squared_euclidean, manhattan, levenshtein
from .test_distances import cosine, transformed_dot_product, jaccard, weighted_jaccard

if sys.version_info.major > 2:
    long = int

import os as _os

class NearestNeighborsCreateTest(unittest.TestCase):
    """
    Unit test class for testing the model create process for various nearest
    neighbors models and methods.
    """

    @classmethod
    def setUpClass(self):

        ## Make data
        np.random.seed(19)
        n, d = 108, 3

        array_features = []
        dict_features = []
        for i in range(n):
            array_features.append(array.array("f", np.random.rand(d)))
            dict_features.append(
                {
                    "alice": np.random.randint(10),
                    "brian": np.random.randint(10),
                    "chris": np.random.randint(10),
                }
            )

        self.refs = tc.SFrame()
        for i in range(d):
            self.refs["X{}".format(i + 1)] = tc.SArray(np.random.rand(n))

        self.label = "label"
        self.refs[self.label] = [str(x) for x in range(n)]
        self.refs["array_ftr"] = array_features
        self.refs["dict_ftr"] = dict_features
        self.refs["str_ftr"] = random_string(n, length=3, num_letters=5)
        self.refs["list_str_ftr"] = random_list_of_str(n, length=3)

    def _test_create(
        self, sf, label, features, distance, method, field=None, value=None
    ):
        """
        Test creation of nearest neighbors models.
        """

        m = tc.nearest_neighbors.create(
            sf, label, features, distance, method, verbose=False
        )
        assert m is not None, "Model creation failed."

        if field is not None:
            ans = m._get(field)
            assert ans == value, "Field '{}' is not correct.".format(field)

    def test_create_default(self):
        """
        Test the default nearest neighbors create configuration.
        """

        ## check auto configurations for the method when features are provided.
        self._test_create(
            self.refs,
            self.label,
            features=["X1", "X2", "X3"],
            distance="auto",
            method="auto",
            field="method",
            value="ball_tree",
        )

        self._test_create(
            self.refs,
            self.label,
            features=["X1", "X2", "X3"],
            distance="euclidean",
            method="auto",
            field="method",
            value="ball_tree",
        )

        ## check auto configurations for distance if features specified.
        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=["X1", "X2", "X3"],
            distance="auto",
            method="brute_force",
            verbose=False,
        )
        self.assertEqual(m.distance, [[["X1", "X2", "X3"], "euclidean", 1.0]])

        ## check auto configurations for distance if features *not* specified.
        ans_dist = [
            [["X1", "X2", "X3"], "euclidean", 1.0],
            [["dict_ftr"], "jaccard", 1.0],
            [["str_ftr"], "levenshtein", 1.0],
            [["array_ftr"], "euclidean", 1.0],
            [["list_str_ftr"], "jaccard", 1.0],
        ]

        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=None,
            distance="auto",
            method="brute_force",
            verbose=False,
        )
        self.assertItemsEqual(m.distance, ans_dist)

        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=None,
            distance=None,
            method="brute_force",
            verbose=False,
        )
        self.assertItemsEqual(m.distance, ans_dist)

        ## check default leaf size for ball tree
        correct_leaf_size = 1000

        self._test_create(
            self.refs,
            self.label,
            features=["array_ftr"],
            distance="euclidean",
            method="ball_tree",
            field="leaf_size",
            value=correct_leaf_size,
        )

        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=["array_ftr"],
            method="ball_tree",
            leaf_size=0,
            verbose=False,
        )
        assert m is not None, "Model creation failed."
        assert m.leaf_size == correct_leaf_size, (
            "Leaf size explicit default" + "failed."
        )

    def test_create_labels(self):
        """
        Test nearest neighbor model creation with different label configurations.
        """

        sf = self.refs[:]
        sf.remove_column(self.label, inplace=True)

        ## String labels are tested everywhere else in this class.

        ## Passing no label should work, with and without listed features
        self._test_create(
            sf,
            label=None,
            features=None,
            distance="auto",
            method="auto",
            field="label",
            value=None,
        )

        self._test_create(
            sf,
            label=None,
            features=["X1", "X2", "X3"],
            distance="euclidean",
            method="auto",
            field="label",
            value=None,
        )

        ## Integer label should work
        sf = sf.add_row_number(column_name="id")

        self._test_create(
            sf,
            label="id",
            features=None,
            distance="auto",
            method="auto",
            field="label",
            value="id",
        )

        m = tc.nearest_neighbors.create(
            sf,
            label="id",
            features=["X1", "X2", "X3"],
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )
        self.assertEqual(set(m.features), set(["X1", "X2", "X3"]))

        ## Float label should fail
        sf["id"] = sf["id"].astype(float)

        with self.assertRaises(TypeError):
            m = tc.nearest_neighbors.create(sf, label="id")

        ## Specified label, included in the features list should drop the label
        #  from the features.
        m = tc.nearest_neighbors.create(
            sf,
            label=None,
            features=["X1", "X2", "__id"],
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )
        self.assertEqual(set(m.features), set(["X1", "X2"]))

        ## If there is only one feature, and it's specified as the label, this
        #  should raise an informative error.
        sf = sf.add_row_number("id_test")

        with self.assertRaises(ToolkitError):
            m = tc.nearest_neighbors.create(sf, label="id_test", features=["id_test"])

    def test_create_methods(self):
        """
        Test different methods of nearest neighbors models.
        """

        methods = {
            "auto": "ball_tree",
            "brute_force": "brute_force",
            "ball_tree": "ball_tree",
            "lsh": "lsh",
        }

        for m, name in methods.items():
            self._test_create(
                self.refs,
                self.label,
                features=["array_ftr"],
                method=m,
                distance="euclidean",
                field="method",
                value=name,
            )

        ## Cosine and transformed_dot_product distances should not work with ball tree
        for dist in [
            "cosine",
            "transformed_dot_product",
            tc.distances.cosine,
            tc.distances.transformed_dot_product,
        ]:

            with self.assertRaises(TypeError):
                tc.nearest_neighbors.create(
                    self.refs,
                    self.label,
                    features=["array_ftr"],
                    distance=dist,
                    method="ball_tree",
                    verbose=False,
                )

        ## Multiple distance components should cause an automatic switch to
        #  brute force, even if ball tree is specified.
        distance_components = [
            [["X1", "X2", "X3"], "euclidean", 1],
            [["array_ftr"], "manhattan", 1],
            [["str_ftr"], "levenshtein", 1],
        ]

        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            method="ball_tree",
            distance=distance_components,
            verbose=False,
        )
        self.assertEqual(m.method, "brute_force")

    def test_kwargs(self):
        """
        Make sure only allowed keyword args are processed. All others should
        raise errors to avoid confusion in downstream errors.
        """
        with self.assertRaises(ToolkitError):
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                feature="array_ftr",  # this is bogus
                method="ball_tree",
                distance="euclidean",
                verbose=False,
            )

    def test_create_dense_distances(self):
        """
        Test various distances in nearest neighbors models.
        """

        dense_dists = {
            "euclidean": tc.distances.euclidean,
            "squared_euclidean": tc.distances.squared_euclidean,
            "gaussian_kernel": tc.distances.gaussian_kernel,
            "manhattan": tc.distances.manhattan,
            "cosine": tc.distances.cosine,
            "transformed_dot_product": tc.distances.transformed_dot_product,
        }

        for dist_name, dist_fn in dense_dists.items():

            ans_dist = [[["array_ftr"], dist_name, 1.0]]

            ## Test the string form of the distance argument
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                features=["array_ftr"],
                distance=dist_name,
                method="brute_force",
                verbose=False,
            )
            self.assertEqual(m.distance, ans_dist)

            ## Test the function form of the distance argument
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                features=["array_ftr"],
                distance=dist_fn,
                method="brute_force",
                verbose=False,
            )
            self.assertEqual(m.distance, ans_dist)

            ## Numeric distances should *not* work with string features
            try:
                tc.nearest_neighbors.create(
                    self.refs, self.label, ["str_ftr"], distance=dist_name
                )
            except ToolkitError as e:
                self.assertTrue(
                    str(e).startswith(
                        "The only distance allowed for string features is 'levenshtein'. "
                        "Please try this distance, or use 'text_analytics.count_ngrams' to "
                        "convert the strings to dictionaries, which permit more distance functions.\n"
                    )
                )

    def test_create_sparse_distances(self):
        """
        Test various distances in nearest neighbors models that operate on
        sparse data, e.g. vectors."""

        sparse_dists = {
            "jaccard": tc.distances.jaccard,
            "weighted_jaccard": tc.distances.weighted_jaccard,
            "cosine": tc.distances.cosine,
            "transformed_dot_product": tc.distances.transformed_dot_product,
        }

        for dist_name, dist_fn in sparse_dists.items():

            ans_dist = [[["dict_ftr"], dist_name, 1.0]]

            ## Test the string form of the distance argument
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                features=["dict_ftr"],
                distance=dist_name,
                method="brute_force",
                verbose=False,
            )
            self.assertEqual(m.distance, ans_dist)

            ## Test the function form of the distance argument
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                features=["dict_ftr"],
                distance=dist_fn,
                method="brute_force",
                verbose=False,
            )
            self.assertEqual(m.distance, ans_dist)

            ## Test the string form of the distance argument with list of str
            ans_dist = [[["list_str_ftr"], dist_name, 1.0]]
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                features=["list_str_ftr"],
                distance=dist_name,
                method="brute_force",
                verbose=False,
            )
            self.assertEqual(m.distance, ans_dist)

        ## Jaccard distances should not work with numeric or string features
        for dist in [
            "jaccard",
            "weighted_jaccard",
            tc.distances.jaccard,
            tc.distances.weighted_jaccard,
        ]:

            try:
                tc.nearest_neighbors.create(
                    self.refs,
                    self.label,
                    features=["array_ftr"],
                    distance=dist,
                    method="brute_force",
                    verbose=False,
                )
            except ToolkitError as e:
                self.assertTrue(
                    str(e).startswith(
                        "Cannot compute jaccard distances with column 'array_ftr'."
                        " Jaccard distances currently can only be computed for"
                        " dictionary and list features.\n"
                    )
                )

            try:
                tc.nearest_neighbors.create(
                    self.refs,
                    self.label,
                    ["str_ftr"],
                    distance=dist_name,
                    verbose=False,
                )
            except ToolkitError as e:
                self.assertTrue(
                    str(e).startswith(
                        "The only distance allowed for string features is 'levenshtein'. "
                        "Please try this distance, or use 'text_analytics.count_ngrams' "
                        "to convert the strings to dictionaries, which permit more distance functions.\n"
                    )
                )

        ## Jacard distance throws TypeError on lists of non-strings
        refs = self.refs.__copy__()
        refs["list_float_ftr"] = refs["array_ftr"].apply(lambda x: list(x), dtype=list)

        # Check autodistance
        with self.assertRaises(TypeError):
            m = tc.nearest_neighbors.create(
                refs, self.label, features=["list_float_ftr"], verbose=False
            )

        # Check user-specified distance
        for distance in ["jaccard", "weighted_jaccard", "euclidean"]:
            with self.assertRaises(TypeError):
                m = tc.nearest_neighbors.create(
                    refs,
                    self.label,
                    features=["list_float_ftr"],
                    distance=distance,
                    method="brute_force",
                    verbose=False,
                )

    def test_create_string_distances(self):
        """
        Test that various string distances can be used to create a nearest
        neighbors model.
        """

        string_dists = {"levenshtein": tc.distances.levenshtein}

        for dist_name, dist_fn in string_dists.items():

            ans_dist = [[["str_ftr"], dist_name, 1.0]]

            ## Test the string form of the distance argument
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                features=["str_ftr"],
                distance=dist_name,
                method="brute_force",
                verbose=False,
            )
            self.assertEqual(m.distance, ans_dist)

            ## Test the function form of the distance argument
            m = tc.nearest_neighbors.create(
                self.refs,
                self.label,
                features=["str_ftr"],
                distance=dist_fn,
                method="brute_force",
                verbose=False,
            )
            self.assertEqual(m.distance, ans_dist)

            ## String distances should not work with numeric or dictionary
            #  features
            try:
                tc.nearest_neighbors.create(
                    self.refs,
                    self.label,
                    features=["dict_ftr"],
                    distance=dist_name,
                    method="brute_force",
                    verbose=False,
                )
            except ToolkitError as e:
                self.assertTrue(
                    str(e).startswith(
                        "Cannot compute {} distance with column 'dict_ftr'.".format(
                            dist_name
                        )
                        + " {} distance can only computed for string features.\n".format(
                            dist_name
                        )
                    )
                )

    def test_create_composite_distances(self):
        """
        Test that nearest neighbors models can be successfully created with
        composite distance parameters.
        """

        distance_components = [
            [["X1", "X2"], "euclidean", 1],
            [
                ["X2", "X3"],
                "manhattan",
                1,
            ],  # note overlap with first component's features
            [["array_ftr"], "manhattan", 1],
            [["str_ftr"], "levenshtein", 1],
        ]

        ## Test that things work correctly in the vanilla case.
        m = tc.nearest_neighbors.create(
            self.refs, self.label, distance=distance_components, verbose=False
        )

        assert m is not None, "Model creation failed."
        self.assertEqual(m.distance, distance_components)
        self.assertEqual(m.num_distance_components, 4)
        self.assertEqual(m.method, "brute_force")
        self.assertEqual(m.num_features, 5)
        self.assertEqual(m.num_unpacked_features, 7)

        ## Make sure the features parameter is ignored if a composite distance
        #  is specified.
        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=["X1", "X2"],
            distance=distance_components,
            verbose=False,
        )

        assert m is not None, "Model creation failed."
        self.assertEqual(m.distance, distance_components)
        self.assertEqual(m.num_distance_components, 4)
        self.assertEqual(m.num_features, 5)
        self.assertEqual(m.num_unpacked_features, 7)

    def test_create_num_variables(self):
        """
        Test vector features, numeric features, and combinations thereof.
        """
        for ftr_list, v in [
            (["X1", "X2", "X3"], 3),
            (["array_ftr"], 3),
            (["dict_ftr"], 3),
        ]:

            self._test_create(
                self.refs,
                self.label,
                features=ftr_list,
                method="auto",
                distance="auto",
                field="num_unpacked_features",
                value=v,
            )

        for ftr_list, v in [
            (["X1", "X2", "X3"], 3),
            (["array_ftr"], 1),
            (["dict_ftr"], 1),
        ]:

            self._test_create(
                self.refs,
                self.label,
                features=ftr_list,
                method="auto",
                distance="auto",
                field="num_features",
                value=v,
            )

    def test_create_mutations(self):
        """
        Test that inputs to the nearest neighbors create function are not
        mutated by the function.
        """
        sf = self.refs[:]
        label = self.label
        ftrs_orig = ["X1", "X2", "X3", "array_ftr"]
        ftrs_copy = ftrs_orig[:]

        m = tc.nearest_neighbors.create(
            sf,
            label=label,
            features=ftrs_copy,
            method="auto",
            distance="auto",
            verbose=False,
        )

        assert_frame_equal(self.refs.to_dataframe(), sf.to_dataframe())
        self.assertEqual(label, self.label)
        self.assertEqual(ftrs_orig, ftrs_copy)

    def test_missing_data(self):
        """
        Check that the create function errors out (correctly) if there are
        missing data of any type in any cell of the input dataset.
        """

        sf = tc.SFrame({"x0": [1, 2, 3], "x1": ["a", "b", "c"]})
        sf["ints"] = [1, 2, None]
        sf["floats"] = [None, 2.2, 3.3]
        sf["strings"] = ["a", None, "c"]
        sf["dicts"] = [{"a": 1}, {"b": 2}, None]
        sf["arrays"] = [
            array.array("f", [1.0, 2.0]),
            array.array("f", [3.0, 4.0]),
            None,
        ]

        for ftr in ["ints", "floats", "strings", "dicts", "arrays"]:
            with self.assertRaises(ToolkitError):
                m = tc.nearest_neighbors.create(sf[["x0", "x1", ftr]], verbose=False)

class NearestNeighborsEdgeCaseTests(unittest.TestCase):
    """"
    Test bad inputs and edge values for inputs.
    """

    @classmethod
    def setUpClass(self):

        ## Make data
        np.random.seed(19)
        n, d = 100, 3
        self.label = "label"

        array_features = []
        for i in range(n):
            array_features.append(array.array("f", np.random.rand(d)))

        self.refs = tc.SFrame()
        self.refs["array"] = array_features
        self.refs[self.label] = [str(x) for x in range(n)]

    def test_empty_data(self):
        """
        Test the errors for empty reference and query datasets.
        """

        ## Useful objects
        sf_empty = tc.SFrame()

        m = tc.nearest_neighbors.create(
            self.refs,
            label=self.label,
            features=None,
            method="brute_force",
            distance="euclidean",
            verbose=False,
        )

        with self.assertRaises(ToolkitError):
            m = tc.nearest_neighbors.create(sf_empty, self.label)

        with self.assertRaises(ToolkitError):
            knn = tc.nearest_neighbors.create(sf_empty, self.label)

    def test_bogus_labels(self):
        """
        Test the errors with labels that don't exist.
        """

        with self.assertRaises(ToolkitError):
            m = tc.nearest_neighbors.create(self.refs, label="fossa")

        m = tc.nearest_neighbors.create(self.refs)

        with self.assertRaises(ValueError):
            m.query(self.refs, label="fossa")

    def test_empty_composite_distance(self):
        """
        Test that empty composite distance parameters are not allowed.
        """

        with self.assertRaises(ValueError):
            m = tc.nearest_neighbors.create(self.refs, distance=[])

    def test_bogus_parameters(self):
        """
        Test toolkit behavior when arguments are inappropriate.
        """

        ## k is out of bounds raises an error
        m = tc.nearest_neighbors.create(self.refs, self.label, method="brute_force")

        for k in [-1, 0, "cat"]:
            with self.assertRaises(ValueError):
                knn = m.query(self.refs, self.label, k=k)

        ## k > n should default to n
        knn = m.query(self.refs, self.label, k=2 * self.refs.num_rows())
        assert (
            knn.num_rows() == self.refs.num_rows() ** 2
        ), "Query with k > n returned the wrong number of rows."

        ## radius out of bounds should raise an error
        for r in [-1, "cat"]:
            with self.assertRaises(ValueError):
                knn = m.query(self.refs, self.label, radius=r)

        # ## Leaf size is out of bounds
        # import ipdb
        # ipdb.set_trace()

        for ls in [-1, 12.3, "fossa"]:
            with self.assertRaises(ToolkitError):
                m = tc.nearest_neighbors.create(
                    self.refs,
                    self.label,
                    features=["array"],
                    method="ball_tree",
                    distance="euclidean",
                    leaf_size=ls,
                )

        ## Leaf size > n should be fine
        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=["array"],
            method="ball_tree",
            distance="euclidean",
            leaf_size=2 * self.refs.num_rows(),
        )
        assert m.leaf_size == 2 * self.refs.num_rows()

        ## Distance component weights are out of bounds
        with self.assertRaises(ValueError):
            m = tc.nearest_neighbors.create(
                self.refs, self.label, distance=[[["array"], "euclidean", -1e-7]]
            )

        with self.assertRaises(ValueError):
            m = tc.nearest_neighbors.create(
                self.refs, self.label, distance=[[["array"], "euclidean", -1]]
            )

        with self.assertRaises(ValueError):
            m = tc.nearest_neighbors.create(
                self.refs, self.label, distance=[[["array"], "euclidean", "a"]]
            )

        with self.assertRaises(ToolkitError):
            m = tc.nearest_neighbors.create(
                self.refs, self.label, distance=[[["array"], "euclidean", 1e15]]
            )

class NearestNeighborsBruteForceAPITest(unittest.TestCase):
    """
    Unit test class for testing model methods once models are trained. Excludes
    the correctness of the query results.
    """

    def setUp(self):

        ## Make data
        np.random.seed(19)
        d = 3  # dimension
        n = 100  # number of reference points

        self.refs = tc.SFrame()
        for i in range(d):
            self.refs.add_column(tc.SArray(np.random.rand(n)), inplace=True)

        self.refs["row_label"] = [str(x) for x in range(n)]

        self.label = "row_label"
        self.features = ["X{}".format(i + 1) for i in range(d)]
        self.unpacked_features = ["X{}".format(i + 1) for i in range(d)]

        ## Create the nearest neighbors model
        self.model = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=None,
            method="brute_force",
            distance="euclidean",
            verbose=False,
        )

        ## Answers
        self.fields_ans = [
            "training_time",
            "label",
            "unpacked_features",
            "features",
            "method",
            "num_examples",
            "num_unpacked_features",
            "num_features",
            "num_distance_components",
            "distance",
            "distance_for_summary_struct",
        ]

        self.default_opts = {
            "leaf_size": {
                u"default_value": 0,
                u"lower_bound": 0,
                u"upper_bound": 2147483647,
                u"description": u"Max number of points in a leaf node of the ball tree",
                u"parameter_type": u"INTEGER",
            },
            "label": {
                u"default_value": u"",
                u"description": u"Name of the reference dataset column with row labels.",
                u"parameter_type": u"STRING",
            },
        }

        self.opts = {"label": self.label}

        self.get_ans = {
            "distance": lambda x: len(x) == 1,
            "training_time": lambda x: x >= 0,
            "label": lambda x: x == self.label,
            "method": lambda x: x == "brute_force",
            "num_examples": lambda x: x == 100,
            "num_features": lambda x: x == 3,
            "num_unpacked_features": lambda x: x == 3,
            "num_distance_components": lambda x: x == 1,
        }

    def test__list_fields(self):
        """
        Check the _list_fields method.
        """
        assert set(self.model._list_fields()) == set(
            self.fields_ans
        ), "List fields failed with {}.".format(self.model._list_fields())

    def test_get(self):
        """
        Check the get method against known answers for each field.
        """
        for field in self.get_ans.keys():
            ans = self.model._get(field)
            assert self.get_ans[field](
                ans
            ), "Get failed in field '{}'. Output: {}".format(field, ans)

        ## Check names of features and unpacked features
        assert set(self.model.features) == set(self.features)
        assert set(self.model.unpacked_features) == set(self.unpacked_features)

    def test_summary(self):
        """
        Check that the summary method returns information.
        """
        try:
            self.model.summary()
        except:
            assert False, "Model summary failed."

        assert self.model.summary() is not ""

    def test_repr(self):
        """
        """
        try:
            ans = str(self.model)
        except:
            assert False, "Model repr failed."

    def test_save_and_load(self):
        """
        Make sure saving and loading retains everything.
        """
        with util.TempDirectory() as f:
            self.model.save(f)
            self.model = tc.load_model(f)

            try:
                self.test__list_fields()
                self.test_get()
            except:
                assert False, "List fields or get failed after save and load."

            try:
                self.model.query(self.refs, self.label, k=3)
            except:
                assert False, "Querying failed after save and load."

            try:
                g = self.model.similarity_graph()
            except:
                assert False, "Similarity failed after save and load."

            try:
                self.test_summary()
                self.test_repr()
            except:
                assert False, "Summaries failed after save and load."

            del self.model

class NearestNeighborsLshAPITest(unittest.TestCase):
    """
    Unit test class for testing model methods once models are trained. Excludes
    the correctness of the query results.
    """

    def setUp(self):

        ## Make data
        np.random.seed(19)
        d = 3  # dimension
        n = 100  # number of reference points

        refs = []
        for i in range(n):
            refs.append(array.array("f", np.random.rand(d)))

        self.refs = tc.SFrame({"features": refs})
        self.refs["row_label"] = [str(x) for x in range(n)]
        self.query = tc.SFrame({"features": refs})
        self.query["row_label"] = [str(x) for x in range(50, n + 50)]

        self.label = "row_label"
        self.features = ["features"]
        self.unpacked_features = ["features[0]", "features[1]", "features[2]"]

        self.opts = {
            "num_tables": 4,
            "num_projections_per_table": 4,
            "label": self.label,
        }

        ## Create the nearest neighbors model
        self.model = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=self.features,
            method="lsh",
            distance="euclidean",
            num_tables=self.opts["num_tables"],
            num_projections_per_table=self.opts["num_projections_per_table"],
        )
        ## Answers
        self.fields_ans = [
            "distance",
            "distance_for_summary_struct",
            "num_distance_components",
            "features",
            "unpacked_features",
            "label",
            "num_tables",
            "num_projections_per_table",
            "method",
            "num_examples",
            "num_unpacked_features",
            "num_features",
            "training_time",
        ]

        self.default_opts = {
            "num_tables": {
                u"default_value": 10,
                u"lower_bound": 1,
                u"upper_bound": 2147483647,
                u"description": u"number of hash tables for LSH",
                u"parameter_type": u"INTEGER",
            },
            "num_projections_per_table": {
                u"default_value": 8,
                u"lower_bound": 1,
                u"upper_bound": 2147483647,
                u"description": u"number of projections in each hash table",
                u"parameter_type": u"INTEGER",
            },
            "label": {
                u"default_value": u"",
                u"description": u"Name of the reference dataset column with row labels.",
                u"parameter_type": u"STRING",
            },
        }

        self.get_ans = {
            "distance": lambda x: len(x) == 1,
            "num_distance_components": lambda x: x == 1,
            "label": lambda x: x == self.label,
            "num_tables": lambda x: x == self.opts["num_tables"],
            "num_projections_per_table": lambda x: x
            == self.opts["num_projections_per_table"],
            "method": lambda x: x == "lsh",
            "num_examples": lambda x: x == 100,
            "num_features": lambda x: x == 1,
            "num_unpacked_features": lambda x: x == 3,
            "training_time": lambda x: x >= 0,
        }

    def test_query(self):
        q = self.model.query(self.query, label=self.label, k=1, verbose=False)
        assert q is not None
        assert q.num_rows() >= self.query.num_rows()
        # all the 1-nearest-neighbor should be the queries themselves (and identical points)
        # so that means a distance of zero
        distances = q["distance"]
        assert len(distances.filter(lambda x: x != 0.0)) == 0

    def test__list_fields(self):
        """
        Check the _list_fields method.
        """
        assert set(self.model._list_fields()) == set(
            self.fields_ans
        ), "List fields failed with {}.".format(self.model._list_fields())

    def test_get(self):
        """
        Check the get method against known answers for each field.
        """
        for field in self.get_ans.keys():
            ans = self.model._get(field)
            assert self.get_ans[field](
                ans
            ), "Get failed in field '{}'. Output: {}".format(field, ans)

        ## Check names of features and unpacked features
        assert set(self.model.features) == set(self.features)
        assert set(self.model.unpacked_features) == set(self.unpacked_features)

    def test_summary(self):
        """
        Check that the summary method returns information.
        """
        try:
            self.model.summary()
        except:
            assert False, "Model summary failed."

        assert self.model.summary() is not ""

    def test_repr(self):
        """
        """
        try:
            ans = str(self.model)
        except:
            assert False, "Model repr failed."

    def test_save_and_load(self):
        """
        Make sure saving and loading retains everything.
        """
        with util.TempDirectory() as f:
            self.model.save(f)
            self.model = tc.load_model(f)

            try:
                print(self.model._list_fields())
                self.test__list_fields()
                self.test_get()
            except:
                assert False, "List fields or get failed after save and load."

            try:
                self.test_summary()
                self.test_repr()
            except:
                assert False, "Summaries failed after save and load."
            del self.model

class NearestNeighborsBallTreeAPITest(unittest.TestCase):
    """
    Unit test class for testing model methods once models are trained. Excludes
    the correctness of the query results.
    """

    def setUp(self):

        ## Make data
        np.random.seed(19)
        d = 3  # dimension
        n = 100  # number of reference points

        refs = []
        for i in range(n):
            refs.append(array.array("f", np.random.rand(d)))

        self.refs = tc.SFrame({"features": refs})
        self.refs["row_label"] = [str(x) for x in range(n)]
        self.query = tc.SFrame({"features": refs})
        self.query["row_label"] = [str(x) for x in range(50, n + 50)]

        self.label = "row_label"
        self.features = ["features"]
        self.unpacked_features = ["features[0]", "features[1]", "features[2]"]

        self.opts = {"leaf_size": 16, "label": self.label}

        ## Create the nearest neighbors model
        self.model = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=self.features,
            method="ball_tree",
            distance="euclidean",
            leaf_size=self.opts["leaf_size"],
            verbose=False,
        )

        ## Answers
        self.fields_ans = [
            "distance",
            "distance_for_summary_struct",
            "num_distance_components",
            "features",
            "unpacked_features",
            "label",
            "leaf_size",
            "method",
            "num_examples",
            "num_unpacked_features",
            "num_features",
            "training_time",
            "tree_depth",
        ]

        self.default_opts = {
            "leaf_size": {
                u"default_value": 0,
                u"lower_bound": 0,
                u"upper_bound": 2147483647,
                u"description": u"Max number of points in a leaf node of the ball tree",
                u"parameter_type": u"INTEGER",
            },
            "label": {
                u"default_value": u"",
                u"description": u"Name of the reference dataset column with row labels.",
                u"parameter_type": u"STRING",
            },
        }

        self.get_ans = {
            "distance": lambda x: len(x) == 1,
            "num_distance_components": lambda x: x == 1,
            "label": lambda x: x == self.label,
            "leaf_size": lambda x: x == self.opts["leaf_size"],
            "method": lambda x: x == "ball_tree",
            "num_examples": lambda x: x == 100,
            "num_features": lambda x: x == 1,
            "num_unpacked_features": lambda x: x == 3,
            "training_time": lambda x: x >= 0,
            "tree_depth": lambda x: x == 4,  # assumes n=100, leaf_size=16
        }

    def test_query(self):
        q = self.model.query(self.query, self.label, k=3, verbose=False)
        assert q is not None

    def test__list_fields(self):
        """
        Check the _list_fields method.
        """
        assert set(self.model._list_fields()) == set(
            self.fields_ans
        ), "List fields failed with {}.".format(self.model._list_fields())

    def test_get(self):
        """
        Check the get method against known answers for each field.
        """
        for field in self.get_ans.keys():
            ans = self.model._get(field)
            assert self.get_ans[field](
                ans
            ), "Get failed in field '{}'. Output: {}".format(field, ans)

        ## Check names of features and unpacked features
        assert set(self.model.features) == set(self.features)
        assert set(self.model.unpacked_features) == set(self.unpacked_features)

    def test_summary(self):
        """
        Check that the summary method returns information.
        """
        try:
            self.model.summary()
        except:
            assert False, "Model summary failed."

        assert self.model.summary() is not ""

    def test_repr(self):
        """
        """
        try:
            ans = str(self.model)
        except:
            assert False, "Model repr failed."

    def test_save_and_load(self):
        """
        Make sure saving and loading retains everything.
        """
        with util.TempDirectory() as f:
            self.model.save(f)
            self.model = tc.load_model(f)

            try:
                print(self.model._list_fields())
                self.test__list_fields()
                self.test_get()
            except:
                assert False, "List fields or get failed after save and load."

            try:
                self.test_summary()
                self.test_repr()
            except:
                assert False, "Summaries failed after save and load."

            del self.model

class GeneralSimilarityGraphTest(unittest.TestCase):
    """
    Tests that apply to all nearest neighbors similarity graph methods,
    regardless of data type.
    """

    @classmethod
    def setUpClass(self):

        np.random.seed(19)
        self.dimension = 3  # dimension
        n = 10  # number of reference points

        self.refs = tc.SFrame(np.random.rand(n, self.dimension))
        self.features = ["X1.{}".format(i) for i in range(self.dimension)]
        self.refs = self.refs.unpack("X1")

        self.label = "id"
        self.refs = self.refs.add_row_number(self.label)
        self.refs[self.label] = self.refs[self.label].astype(str) + "a"

        df_refs = self.refs.to_dataframe().drop(self.label, axis=1)
        self.answer_dists = scipy_dist(df_refs, df_refs, "euclidean")

    def test_neighborhood_constraints(self):
        """
        Test various combinations of the k and radius constraints.
        """
        m = tc.nearest_neighbors.create(
            self.refs,
            features=self.features,
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )

        knn = m.similarity_graph(
            k=None,
            radius=None,
            include_self_edges=True,
            output_type="SFrame",
            verbose=False,
        )
        assert_frame_equal(self.answer_dists.to_dataframe(), knn.to_dataframe())

        ## k only, no radius
        knn = m.similarity_graph(
            k=3,
            radius=None,
            include_self_edges=True,
            output_type="SFrame",
            verbose=False,
        )
        ans = self.answer_dists[self.answer_dists["rank"] <= 3]
        assert_frame_equal(ans.to_dataframe(), knn.to_dataframe())

        ## radius only, no k
        knn = m.similarity_graph(
            k=None,
            radius=0.4,
            include_self_edges=True,
            output_type="SFrame",
            verbose=False,
        )
        ans = self.answer_dists[self.answer_dists["distance"] <= 0.4]
        assert_frame_equal(ans.to_dataframe(), knn.to_dataframe())

        ## Both radius and k
        knn = m.similarity_graph(
            k=3,
            radius=0.4,
            include_self_edges=True,
            output_type="SFrame",
            verbose=False,
        )
        ans = self.answer_dists[self.answer_dists["rank"] <= 3]
        ans = ans[ans["distance"] <= 0.4]
        assert_frame_equal(ans.to_dataframe(), knn.to_dataframe())

    def test_self_edges(self):
        """
        Check that the 'include_self_edges' flag performs as expected, with and
        without reference row labels.
        """

        ## Without row labels
        m = tc.nearest_neighbors.create(
            self.refs,
            features=self.features,
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )

        knn = m.similarity_graph(
            k=None,
            radius=None,
            include_self_edges=True,
            output_type="SFrame",
            verbose=False,
        )
        assert_frame_equal(self.answer_dists.to_dataframe(), knn.to_dataframe())

        knn2 = m.similarity_graph(
            k=None,
            radius=None,
            include_self_edges=False,
            output_type="SFrame",
            verbose=False,
        )
        mask = self.answer_dists["query_label"] != self.answer_dists["reference_label"]
        ans = self.answer_dists[mask]
        ans["rank"] = ans["rank"] - 1
        assert_frame_equal(ans.to_dataframe(), knn2.to_dataframe())

        ## With string row labels
        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=None,
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )

        knn = m.similarity_graph(
            k=None,
            radius=None,
            output_type="SFrame",
            include_self_edges=True,
            verbose=False,
        )

        ans = copy.copy(self.answer_dists)
        ans["query_label"] = ans["query_label"].astype(str) + "a"
        ans["reference_label"] = ans["reference_label"].astype(str) + "a"
        assert_frame_equal(ans.to_dataframe(), knn.to_dataframe())

        knn2 = m.similarity_graph(
            k=None,
            radius=None,
            include_self_edges=False,
            output_type="SFrame",
            verbose=False,
        )
        mask = self.answer_dists["query_label"] != self.answer_dists["reference_label"]
        ans = self.answer_dists[mask]
        ans["rank"] = ans["rank"] - 1
        ans["query_label"] = ans["query_label"].astype(str) + "a"
        ans["reference_label"] = ans["reference_label"].astype(str) + "a"

        assert_frame_equal(ans.to_dataframe(), knn2.to_dataframe())

    def test_output_type(self):
        """
        Check that the results can be returned as either an SFrame or an SGraph
        and that the results match in both of these forms.
        """
        m = tc.nearest_neighbors.create(
            self.refs,
            features=self.features,
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )

        knn = m.similarity_graph(
            k=None,
            radius=None,
            include_self_edges=False,
            output_type="SFrame",
            verbose=False,
        )

        sg = m.similarity_graph(
            k=None,
            radius=None,
            include_self_edges=False,
            output_type="SGraph",
            verbose=False,
        )

        sg_edges = copy.copy(sg.edges)
        sg_edges = sg_edges.rename(
            {"__src_id": "query_label", "__dst_id": "reference_label"}, inplace=True
        )
        sg_edges = sg_edges.sort(["query_label", "distance"])

        assert_frame_equal(sg_edges.to_dataframe(), knn.to_dataframe())

    def test_other_methods(self):
        """
        Similarity graph should also work on ball_tree and lsh
        """

        ## Ball tree
        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=None,
            distance="euclidean",
            method="ball_tree",
            verbose=False,
        )

        knn = m.similarity_graph(k=5, radius=None, verbose=False)

        ## LSH
        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            features=None,
            distance="euclidean",
            method="lsh",
            verbose=False,
        )

        knn = m.similarity_graph(k=5, radius=None, verbose=False)

class GeneralQueryTest(unittest.TestCase):
    """
    Tests that apply to all nearest neighbors model queries, regardless of data
    type.
    """

    @classmethod
    def setUpClass(self):

        np.random.seed(19)
        p = 3  # dimension
        n = 10  # number of reference points
        self.n_query = 2  # number of query points

        self.refs = tc.SFrame(np.random.rand(n, p))
        self.refs = self.refs.unpack("X1")

        self.label = "id"
        self.refs = self.refs.add_row_number(self.label)
        self.queries = self.refs[0 : self.n_query]

        df_refs = self.refs.to_dataframe().drop(self.label, axis=1)
        df_queries = self.queries.to_dataframe().drop(self.label, axis=1)
        self.answer_dists = scipy_dist(df_queries, df_refs, "euclidean")

    def test_neighborhood_constraints(self):
        """
        Test various combinations of the k and radius constraints.
        """

        ## No constraints
        m = tc.nearest_neighbors.create(
            self.refs,
            self.label,
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )

        knn = m.query(self.queries, k=None, radius=None, verbose=False)

        assert_frame_equal(self.answer_dists.to_dataframe(), knn.to_dataframe())

        ## k only, no radius
        knn = m.query(self.queries, k=3, radius=None, verbose=False)
        ans = self.answer_dists[self.answer_dists["rank"] <= 3]
        assert_frame_equal(ans.to_dataframe(), knn.to_dataframe())

        ## radius only, no k
        knn = m.query(self.queries, k=None, radius=0.4, verbose=False)
        ans = self.answer_dists[self.answer_dists["distance"] <= 0.4]
        assert_frame_equal(ans.to_dataframe(), knn.to_dataframe())

        ## Both radius and k
        knn = m.query(self.queries, k=3, radius=0.4, verbose=False)
        ans = self.answer_dists[self.answer_dists["rank"] <= 3]
        ans = ans[ans["distance"] <= 0.4]
        assert_frame_equal(ans.to_dataframe(), knn.to_dataframe())

    def test_labels(self):
        """
        Test query accuracy for various configurations of row labels.
        """
        sfq = self.queries[:]
        sfq.remove_column("id", inplace=True)
        k = 3

        m = tc.nearest_neighbors.create(
            self.refs,
            label=self.label,
            features=None,
            distance="euclidean",
            method="brute_force",
            verbose=False,
        )
        knn_correct = m.query(self.queries, self.label, k=k, verbose=False)

        ## No label should work fine
        knn = m.query(sfq, label=None, k=k, verbose=False)
        self.assertTrue(knn["query_label"].dtype is int)
        assert_frame_equal(knn_correct.to_dataframe(), knn.to_dataframe())

        ## Integer label should work fine
        sfq = sfq.add_row_number(column_name="id")
        knn = m.query(sfq, label="id", k=k, verbose=False)
        self.assertTrue(knn["query_label"].dtype is int)
        assert_frame_equal(knn_correct.to_dataframe(), knn.to_dataframe())

        ## Float label should fail
        sfq["id"] = sfq["id"].astype(float)
        with self.assertRaises(TypeError):
            knn = m.query(sfq, label="id", k=k)

class NearestNeighborsNumericQueryTest(unittest.TestCase):
    """
    Unit test class for checking correctness of exact model queries.
    """

    @classmethod
    def setUpClass(self):

        ## Make data
        np.random.seed(19)
        p = 3  # dimension
        n = 10  # number of reference points
        self.n_query = 2  # number of query points

        self.refs = tc.SFrame(np.random.rand(n, p))
        self.refs = self.refs.unpack("X1")

        self.label = "id"
        self.refs = self.refs.add_row_number(self.label)
        self.queries = self.refs[0 : self.n_query]

        ## Answer items common to all tests
        self.r = self.refs.to_dataframe().drop(self.label, axis=1)
        self.q = self.queries.to_dataframe().drop(self.label, axis=1)

    def _test_query(
        self,
        answer,
        sf_ref,
        sf_query,
        label,
        features,
        distance,
        method,
        k=None,
        radius=None,
    ):
        """
        Test the accuracy of exact queries against python brute force solution,
        from Scipy.
        """

        ## Construct nearest neighbors model and get query results
        m = tc.nearest_neighbors.create(
            sf_ref, label, features, distance, method, verbose=False
        )
        knn = m.query(sf_query, label, k=k, radius=radius, verbose=False)

        ## Trim the answers to the k and radius parameters
        if k is not None:
            answer = answer[answer["rank"] <= k]

        if radius is not None:
            answer = answer[answer["distance"] <= radius]

        ## Test data frame equality
        assert_frame_equal(answer.to_dataframe(), knn.to_dataframe())

    def test_query_distances(self):
        """
        Test query accuracy for various distances.
        """

        idx_row = np.array([[x] for x in range(self.n_query)])

        ## Euclidean and cosine distance
        for dist in ["euclidean", "cosine"]:
            answer = scipy_dist(self.q, self.r, dist)
            self._test_query(
                answer,
                self.refs,
                self.queries,
                self.label,
                features=None,
                distance=dist,
                method="brute_force",
            )

        ## Squared euclidean distances
        answer = scipy_dist(self.q, self.r, "sqeuclidean")
        self._test_query(
            answer,
            self.refs,
            self.queries,
            self.label,
            features=None,
            distance="squared_euclidean",
            method="brute_force",
        )

        ## Manhattan distance
        answer = scipy_dist(self.q, self.r, "cityblock")
        self._test_query(
            answer,
            self.refs,
            self.queries,
            self.label,
            features=None,
            distance="manhattan",
            method="brute_force",
        )

        ## Auto distance (brute force, dense features)
        answer = scipy_dist(self.q, self.r, "euclidean")
        self._test_query(
            answer,
            self.refs,
            self.queries,
            self.label,
            features=None,
            distance="auto",
            method="brute_force",
        )

        ## Transformed dot product distance
        D = self.q.dot(self.r.T)
        D = np.log(1 + np.exp(-1 * D))

        D = D.values
        n_query, n = D.shape

        idx_col = np.argsort(D, axis=1)
        idx_row = np.array([[x] for x in range(n_query)])
        query_labels = list(np.repeat(range(n_query), n))
        ranks = np.tile(range(1, n + 1), n_query)

        answer = tc.SFrame(
            {
                "query_label": query_labels,
                "reference_label": idx_col.flatten(),
                "distance": D[idx_row, idx_col].flatten(),
                "rank": ranks,
            }
        )

        answer.swap_columns("distance", "query_label", inplace=True)
        answer.swap_columns("distance", "reference_label", inplace=True)
        answer.swap_columns("distance", "rank", inplace=True)

        self._test_query(
            answer,
            self.refs,
            self.queries,
            self.label,
            features=None,
            distance="transformed_dot_product",
            method="brute_force",
        )

    def test_query_methods(self):
        """
        Test query accuracy for various nearest neighbor methods.
        """
        answer = scipy_dist(self.q, self.r, "euclidean")

        for method in ["auto", "brute_force", "ball_tree"]:
            self._test_query(
                answer,
                self.refs,
                self.queries,
                self.label,
                features=None,
                distance="euclidean",
                method=method,
            )

    def test_blockwise_brute_force(self):
        """
        Test query completeness and accuracy for brute force queries, with the
        block-wise brute force implementation.
        """

        ## On debug, the following dimensions should produce 121 blocks of
        #  reference data, and 2 blocks of query data.
        n, d = 1927, 1000
        n_query = 21  # more than 20 queries is necessary for block-wise brute force
        k = 5

        sf = tc.SFrame(np.random.rand(n, d))

        m = tc.nearest_neighbors.create(
            sf, method="brute_force", distance="euclidean", verbose=False
        )

        sf_query = tc.SFrame(np.random.rand(n_query, d))

        knn = m.query(sf_query, verbose=False)  # blockwise brute force query
        knn2 = m.query(sf_query[:10], verbose=False)  # pairwise brute force query

        self.assertEqual(knn.num_rows(), 21 * 5)
        assert_frame_equal(knn[: 10 * 5].to_dataframe(), knn2.to_dataframe())

    def test_similarity_graph(self):
        """
        Test accuracy and completeness of the similarity graph method. The
        numeric version also tests the blockwise similarity graph.
        """

        ### Blockwise similarity graph - euclidean distance
        ### -----------------------------------------------
        n, d = 500, 10
        sf = tc.SFrame(np.random.rand(n, d))

        m = tc.nearest_neighbors.create(
            sf, method="brute_force", distance="euclidean", verbose=False
        )

        knn = m.query(sf[:10], k=3, verbose=False)
        knn_graph = m.similarity_graph(k=2, output_type="SFrame", verbose=False)

        ## Basic metadata about the output
        self.assertEqual(knn_graph.num_rows(), 1000)
        self.assertEqual(knn_graph["rank"].max(), 2)
        self.assertGreaterEqual(knn_graph["distance"].min(), 0.0)

        ## No self-edges
        label_diff = knn_graph["query_label"] - knn_graph["reference_label"]
        self.assertEqual(sum(label_diff == 0), 0)

        ## Exact match for query results with different
        knn = knn[knn["rank"] > 1]
        knn["rank"] = knn["rank"] - 1
        assert_frame_equal(knn.to_dataframe(), knn_graph[:20].to_dataframe())

        ### Tweaking query-time parameters
        ### ------------------------------
        knn_graph2 = m.similarity_graph(k=5, output_type="SFrame", verbose=False)
        self.assertEqual(knn_graph2.num_rows(), 2500)
        self.assertEqual(knn_graph2["rank"].max(), 5)
        self.assertGreaterEqual(knn_graph2["distance"].min(), 0.0)

        knn_graph3 = m.similarity_graph(
            k=None, radius=0.5, output_type="SFrame", verbose=False
        )
        self.assertLessEqual(knn_graph3["distance"].max(), 0.5)

        knn_graph4 = m.similarity_graph(
            k=2, radius=0.5, output_type="SFrame", verbose=False
        )
        self.assertEqual(knn_graph4["rank"].max(), 2)
        self.assertLessEqual(knn_graph4["distance"].max(), 0.5)

        ### Pairwise similarity graph - manhattan distance
        ### ----------------------------------------------
        n, d = 500, 10
        sf = tc.SFrame(np.random.rand(n, d))

        m = tc.nearest_neighbors.create(
            sf, method="brute_force", distance="manhattan", verbose=False
        )

        knn = m.query(sf[:10], k=3, verbose=False)
        knn_graph = m.similarity_graph(k=2, output_type="SFrame", verbose=False)

        ## Basic metadata about the output
        self.assertEqual(knn_graph.num_rows(), 1000)
        self.assertEqual(knn_graph["rank"].max(), 2)
        self.assertGreaterEqual(knn_graph["distance"].min(), 0.0)

        ## No self-edges
        label_diff = knn_graph["query_label"] - knn_graph["reference_label"]
        self.assertEqual(sum(label_diff == 0), 0)

        ## Exact match for query results with different
        knn = knn[knn["rank"] > 1]
        knn["rank"] = knn["rank"] - 1
        assert_frame_equal(knn.to_dataframe(), knn_graph[:20].to_dataframe())

class NearestNeighborsSparseQueryTest(unittest.TestCase):
    """
    Unit test class for checking correctness of exact model queries.
    """

    @classmethod
    def setUpClass(self):
        n = 10
        n_query = 2

        # Generate sparse data: an SFrame with a column of type dict
        self.field = "docs"
        self.refs = tc.SFrame()
        self.refs[self.field] = [random_dict(5, 3) for i in range(n)]
        self.k = 3  # number of neighbors to return
        self.radius = 1.0  # radius to use
        self.label = "row_label"
        self.refs[self.label] = [str(x) for x in range(n)]

    def _test_query(
        self, sf_ref, sf_query, label, features, distance, method, k=5, radius=1.0
    ):
        """
        Test the accuracy of exact queries against hand-coded solution above.
        """

        ## Construct nearest neighbors model and get query results
        m = tc.nearest_neighbors.create(
            sf_ref, label, features, distance, method, verbose=False
        )

        knn = m.query(sf_query, label, k=k, radius=radius)

        # TODO: Speed up this test
        for row in knn:
            q = row["query_label"]
            r = row["reference_label"]
            query = sf_ref[int(q)][self.field]
            ref = sf_ref[int(r)][self.field]

            score = row["distance"]

            if distance == "cosine":
                ans = cosine(query, ref)
            elif distance == "dot_product":
                ans = dot_product(query, ref)
            elif distance == "transformed_dot_product":
                ans = transformed_dot_product(query, ref)
            elif distance == "jaccard":
                ans = jaccard(query, ref)
            elif distance == "weighted_jaccard":
                ans = weighted_jaccard(query, ref)
            elif distance == "euclidean":
                ans = euclidean(query, ref)
            elif distance == "squared_euclidean":
                ans = squared_euclidean(query, ref)
            elif distance == "manhattan":
                ans = manhattan(query, ref)
            else:
                raise RuntimeError("Unknown distance")
            self.assertAlmostEqual(score, ans)

    def test_query_distances(self):
        """
        Test query accuracy for various distances.
        """
        for dist in [
            "euclidean",
            "squared_euclidean",
            "manhattan",
            "cosine",
            "transformed_dot_product",
            "jaccard",
            "weighted_jaccard",
        ]:

            self._test_query(
                self.refs,
                self.refs,
                self.label,
                features=None,
                distance=dist,
                method="brute_force",
            )
            self._test_query(
                self.refs,
                self.refs,
                self.label,
                features=None,
                distance=dist,
                method="brute_force",
                k=self.k,
            )
            self._test_query(
                self.refs,
                self.refs,
                self.label,
                features=None,
                distance=dist,
                method="brute_force",
                k=self.k,
                radius=self.radius,
            )

    def test_query_methods(self):
        """
        Test query accuracy for various nearest neighbor methods.
        """
        for method in ["auto", "ball_tree", "brute_force"]:
            self._test_query(
                self.refs,
                self.refs,
                self.label,
                features=None,
                distance="euclidean",
                method=method,
            )

    def test_similarity_graph(self):
        """
        Test accuracy and completeness of the similarity graph method for sparse
        data. In this data context, this is identical to querying with the
        reference dataset.
        """
        n = 30
        sf = tc.SFrame()
        sf["docs"] = [random_dict(5, 3) for i in range(n)]
        sf = sf.add_row_number("id")

        m = tc.nearest_neighbors.create(
            sf, label="id", features=None, distance="euclidean", method="brute_force"
        )

        knn = m.query(sf, k=3, verbose=False)
        knn_graph = m.similarity_graph(k=2, output_type="SFrame", verbose=False)

        ## Basic metadata about the output
        self.assertEqual(knn_graph.num_rows(), 60)
        self.assertEqual(knn_graph["rank"].max(), 2)
        self.assertGreaterEqual(knn_graph["distance"].min(), 0.0)

        ## No self edges
        label_diff = knn_graph["query_label"] - knn_graph["reference_label"]
        self.assertEqual(sum(label_diff == 0), 0)

        ## Exact match to query method, adjusting for no self-edges. NOTE: for
        #  this type of data and distance, there are many ties, so the reference
        #  labels won't necessarily match across the two methods. Only check the
        #  three columns that must match exactly.
        knn = knn[knn["rank"] > 1]
        knn["rank"] = knn["rank"] - 1
        test_ftrs = ["query_label", "distance", "rank"]

        assert_frame_equal(
            knn[test_ftrs].to_dataframe(), knn_graph[test_ftrs].to_dataframe()
        )

class NearestNeighborsStringQueryTest(unittest.TestCase):
    """
    Unit test class for checking correctness of string distances in nearest
    neighbors queries.
    """

    @classmethod
    def setUpClass(self):
        np.random.seed(213)
        n = 5
        word_length = 3
        alphabet_size = 5
        self.label = "id"
        self.refs = tc.SFrame({"X1": random_string(n, word_length, alphabet_size)})
        self.refs = self.refs.add_row_number(self.label)

    def _test_query(self, sf_ref, sf_query, features, distance, method):
        """
        Test the accuracy of string queries against the local python function.
        """

        ## Get the toolkit answer
        m = tc.nearest_neighbors.create(
            sf_ref,
            label=self.label,
            features=features,
            distance=distance,
            method=method,
            verbose=False,
        )
        knn = m.query(sf_query, verbose=False)

        ## Compute the answer from scratch
        knn = knn.join(self.refs, on={"query_label": "id"}, how="left")
        knn = knn.join(self.refs, on={"reference_label": "id"}, how="left")

        if distance == "levenshtein":
            knn["test_dist"] = knn.apply(lambda x: levenshtein(x["X1"], x["X1.1"]))
        else:
            raise ValueError("Distance not found in string query test.")

        self.assertAlmostEqual(sum(knn["distance"] - knn["test_dist"]), 0)

    def test_query_distances(self):
        """
        Test query accuracy for various dista0nces. As of v1.1, only levenshtein
        distance is implemented.
        """
        self._test_query(
            self.refs,
            self.refs,
            features=None,
            distance="levenshtein",
            method="brute_force",
        )

    def test_similarity_graph(self):
        """
        Test accuracy and completeness of the similarity graph method for sparse
        data. In this data context, this is identical to querying with the
        reference dataset.
        """
        n = 30
        sf = tc.SFrame({"X1": random_string(n, length=3, num_letters=5)})
        sf = sf.add_row_number("id")

        m = tc.nearest_neighbors.create(
            sf, label="id", features=None, distance="levenshtein", method="brute_force"
        )

        knn = m.query(sf, k=3, verbose=False)
        knn_graph = m.similarity_graph(k=2, output_type="SFrame", verbose=False)

        ## Basic metadata about the output
        self.assertEqual(knn_graph.num_rows(), 60)
        self.assertEqual(knn_graph["rank"].max(), 2)
        self.assertGreaterEqual(knn_graph["distance"].min(), 0.0)

        ## No self edges
        label_diff = knn_graph["query_label"] - knn_graph["reference_label"]
        self.assertEqual(sum(label_diff == 0), 0)

        ## Exact match to query method, adjusting for no self-edges. NOTE: for
        #  this type of data and distance, there are many ties, so the reference
        #  labels won't necessarily match across the two methods. Only check the
        #  three columns that must match exactly.
        knn = knn[knn["rank"] > 1]
        knn["rank"] = knn["rank"] - 1
        test_ftrs = ["query_label", "distance", "rank"]

        assert_frame_equal(
            knn[test_ftrs].to_dataframe(), knn_graph[test_ftrs].to_dataframe()
        )

    def test_missing_queries(self):
        """
        Check that missing string queries are correctly imputed to be empty
        strings.
        """
        sf = tc.SFrame({"x0": ["a", "b"], "x1": ["d", "e"]})

        sf_query = tc.SFrame(
            {"x0": ["a", None, "b", None], "x1": ["b", "c", None, None]}
        )

        m = tc.nearest_neighbors.create(sf, verbose=False)
        knn = m.query(sf_query, k=None, radius=None)

        answer = tc.SFrame(
            {
                "query_label": [0, 0, 1, 1, 2, 2, 3, 3],
                "distance": [1.0, 2.0, 2.0, 2.0, 1.0, 2.0, 2.0, 2.0],
            }
        )

        assert_frame_equal(
            knn[["distance", "query_label"]].to_dataframe(), answer.to_dataframe()
        )

class NearestNeighborsCompositeQueryTest(unittest.TestCase):
    """
    Unit test class for checking correctness of composite distances in nearest
    neighbors queries.
    """

    @classmethod
    def setUpClass(self):
        n, d = 5, 3
        self.refs = tc.SFrame()

        for i in range(d):
            self.refs.add_column(tc.SArray(np.random.rand(n)), inplace=True)

        # self.refs['address'] = random_string(n, length=3, num_letters=5)
        # self.refs['address_dict'] = tc.text_analytics.count_ngrams(
        #     self.refs['address'], n=2, method='character', to_lower=False,
        #     ignore_space=True)

        self.label = "id"
        self.refs = self.refs.add_row_number(self.label)

    def _test_query(self, composite_params):
        """
        Test query accuracy using arbitrary composite distance inputs.
        """
        ## Get the toolkit answer
        m = tc.nearest_neighbors.create(
            self.refs, label=self.label, distance=composite_params, verbose=False
        )
        knn = m.query(self.refs, verbose=False)

        ## Compute the answer from scratch
        dist_ans = []
        for row in knn:
            query = self.refs[row["query_label"]]
            ref = self.refs[row["reference_label"]]
            dist_row = tc.distances.compute_composite_distance(
                composite_params, query, ref
            )
            dist_ans.append(dist_row)

        knn["test_dist"] = dist_ans
        self.assertAlmostEqual(sum(knn["distance"] - knn["test_dist"]), 0)

    def test_composite_queries(self):
        """
        Test query accurateness for a variety of composite distance
        configurations.
        """

        ## Test accuracy over overlapping feature sets
        distance_components = [
            [["X1", "X2"], "euclidean", 1],
            [["X2", "X3"], "manhattan", 1],
        ]
        self._test_query(distance_components)

        ## Test accuracy of different weights
        distance_components = [
            [["X1", "X2"], "euclidean", 2],
            [["X2", "X3"], "manhattan", 3.4],
        ]
        self._test_query(distance_components)

    def test_similarity_graph(self):
        """
        Test accuracy and completeness of the similarity graph method for sparse
        data. In this data context, this is identical to querying with the
        reference dataset.
        """
        n, d = 30, 3
        sf = tc.SFrame(np.random.random((n, d)))
        sf = sf.unpack("X1", column_name_prefix="")
        sf = sf.add_row_number("id")

        my_dist = [[["0", "1"], "euclidean", 1], [["1", "2"], "manhattan", 1]]

        m = tc.nearest_neighbors.create(
            sf, label="id", features=None, distance=my_dist, method="brute_force"
        )

        knn = m.query(sf, k=3, verbose=False)
        knn_graph = m.similarity_graph(k=2, output_type="SFrame", verbose=False)

        ## Basic metadata about the output
        self.assertEqual(knn_graph.num_rows(), 60)
        self.assertEqual(knn_graph["rank"].max(), 2)
        self.assertGreaterEqual(knn_graph["distance"].min(), 0.0)

        ## No self edges
        label_diff = knn_graph["query_label"] - knn_graph["reference_label"]
        self.assertEqual(sum(label_diff == 0), 0)

        ## Exact match to query method, adjusting for no self-edges.
        knn = knn[knn["rank"] > 1]
        knn["rank"] = knn["rank"] - 1
        test_ftrs = ["query_label", "distance", "rank"]

        assert_frame_equal(knn.to_dataframe(), knn_graph.to_dataframe())

class ValidateListUtilityTest(unittest.TestCase):
    """
    Unit test class for checking correctness of utility function.
    """

    def setUp(self):
        self.check_for_numeric_fixed_length_lists = partial(
            _validate_lists,
            allowed_types=[int, float, long],
            require_equal_length=True,
            require_same_type=True,
        )

        self.check_for_str_lists = partial(
            _validate_lists,
            allowed_types=[str],
            require_equal_length=False,
            require_same_type=True,
        )

    def test_str_cases(self):
        lst = [["a", "b"], ["a", "b", "c"]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_str_lists(lst)
        self.assertTrue(observed)

        lst = [[1, 2, 3], [1, 2, 3, 4]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_str_lists(lst)
        self.assertFalse(observed)

        lst = [["a", 2, 3], [1, 2, 3, 4]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_str_lists(lst)
        self.assertFalse(observed)

    def test_true_numeric_cases(self):

        # Same length, all floats
        lst = [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertTrue(observed)

        # Same length, all ints
        lst = [[1, 2, 3], [2, 3, 4]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertTrue(observed)

        # Only checks the first 10
        lst = [
            [1, 2],
            [2, 3],
            [1, 2],
            [1, 2],
            [2, 3],
            [1, 2],
            [1, 2],
            [2, 3],
            [1, 2],
            [1, 2],
            [2, 3, 4],
            [1, 2, 3],
            [2, 3, 4, 5, 6],
        ]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertTrue(observed)

        # Empty assumed to be same length, all numeric.
        lst = []
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertTrue(observed)

    def test_false_numeric_cases(self):

        # Different length
        lst = [[1, 2, 3], [1, 2, 3, 4]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertFalse(observed)

        # Second list contains different types
        lst = [[1, 2, 3], [1.0, 2, 3]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertFalse(observed)

        # Second list is different typed than first
        lst = [[1, 2, 3], [1.0, 2.0, 3.0]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertFalse(observed)

        # Different length arrays
        lst = [[], [1.0]]
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertFalse(observed)

        # Strings are not numeric
        str_list = random_list_of_str(10, length=5)
        lst = tc.SArray(lst, dtype=list)
        observed = self.check_for_numeric_fixed_length_lists(lst)
        self.assertFalse(observed)

    def test_bad_cases(self):

        lst = [{"a": 3}, {"b": 5}]
        lst = tc.SArray(lst)
        with self.assertRaises(ValueError):
            observed = self.check_for_numeric_fixed_length_lists(lst)

        lst = [1, 2]
        lst = tc.SArray(lst)
        with self.assertRaises(ValueError):
            observed = self.check_for_numeric_fixed_length_lists(lst)

## Generate sparse data: an SFrame with a column of type dict
def random_dict(num_elements=10, max_count=20):
    words = string.ascii_lowercase
    d = {}
    for j in range(num_elements):
        w = words[random.randint(0, len(words) - 1)]
        d[w] = random.randint(0, max_count)
    return d

def random_string(number, length, num_letters):
    """
    Generate a list of random strings of lower case letters.
    """
    result = []
    letters = string.ascii_letters[:num_letters]

    for i in range(number):
        word = []
        for j in range(length):
            word.append(random.choice(letters))
        result.append("".join(word))

    return result

def random_list_of_str(number, length):
    """
    Generate a list of random lists of strings.
    """
    results = []
    words = random_string(20, 5, 10)

    for i in range(number):
        result = [random.choice(words) for i in range(length)]
        results.append(result)
    return results

def scipy_dist(q, r, dist):
    import scipy.spatial.distance as spd

    n = len(r)
    n_query = len(q)

    D = spd.cdist(q, r, dist)
    idx_col = np.argsort(D, axis=1)
    idx_row = np.array([[x] for x in range(n_query)])
    query_labels = list(np.repeat(range(n_query), n))
    ranks = np.tile(range(1, n + 1), n_query)

    answer = tc.SFrame(
        {
            "query_label": query_labels,
            "reference_label": idx_col.flatten(),
            "distance": D[idx_row, idx_col].flatten(),
            "rank": ranks,
        }
    )

    answer.swap_columns("distance", "query_label", inplace=True)
    answer.swap_columns("distance", "reference_label", inplace=True)
    answer.swap_columns("distance", "rank", inplace=True)

    return answer

if __name__ == "__main__":

    # Check if we are supposed to connect to another server
    for i, v in enumerate(sys.argv):
        if v.startswith("ipc://"):
            tc._stop()
            tc._launch(v)

            # The rest of the arguments need to get passed through to
            # the unittest module
            del sys.argv[i]
            break

    unittest.main()


#!/usr/bin/env pypy3
from sys import stdin, exit
from copy import *
from heapq import *
from collections import *
from itertools import *
from functools import *
from math import *
import os
import sys
import requests

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from init import get_session

def submit(part: int, ans):
    res = input(f"Submit answer {ans}? [Y/n]")
    if 'N' in res or 'n' in res: return
    day = int(os.path.basename(os.path.dirname(os.path.realpath(__file__))))
    r = requests.post(f"https://adventofcode.com/2022/day/{day}/answer", data={"level": part, "answer": ans}, cookies={"session": get_session()})
    print(r.text)

with open("input.txt", "r") as f:
    data = f.read().split('\n')[:-1]

def part1():
    ans = 0
    for i, line in enumerate(data):
        for j, val in enumerate(line):
            for x in range(0, i):
                if data[x][j] >= val: break
            else:
                ans += 1
                continue
            for x in range(i + 1, len(data)):
                if data[x][j] >= val: break
            else:
                ans += 1
                continue
            for y in range(0, j):
                if data[i][y] >= val: break
            else:
                ans += 1
                continue
            for y in range(j + 1, len(line)):
                if data[i][y] >= val: break
            else:
                ans += 1
                continue
    submit(1, ans)

def part2():
    ans = 0
    for i, line in enumerate(data):
        for j, val in enumerate(line):
            score = [0, 0, 0, 0]
            for x in range(i - 1, -1, -1):
                score[0] += 1
                if data[x][j] >= val:
                    break
            for x in range(i + 1, len(data)):
                score[1] += 1
                if data[x][j] >= val:
                    break
            for y in range(j - 1, -1, -1):
                score[2] += 1
                if data[i][y] >= val:
                    break
            for y in range(j + 1, len(line)):
                score[3] += 1
                if data[i][y] >= val:
                    break
            ans = max(ans, reduce(lambda x, y: x * y, score))
    submit(2, ans)

if __name__ == "__main__":
    part1()
    part2()


import os
import cv2
import time
import numpy as np
import pandas as pd
import gc

from numpy import nan
from numpy import genfromtxt

# flip frame
def horizontal_flip(src):
	return cv2.flip(src,-1)

# starts and stops the video to find the sychronization point.
def find_sync(video,bindings=0):

	buttons = {}
	if(bindings == 0):
		buttons["pause"] = ord('o')
		buttons["reverse"] = ord('k')
		buttons["quit"] = ord('q')
	if(bindings == 1):
		buttons["pause"] = ord('p')
		buttons["reverse"] = ord('l')
		buttons["quit"] = ord('w')

	vs = cv2.VideoCapture(vidPath + video)

	num_frames = vs.get(cv2.CAP_PROP_FRAME_COUNT)
	FPS = vs.get(cv2.CAP_PROP_FPS)

	video_label = np.zeros(int(num_frames))
	current_frame = 0
	video_frame = []

	# FILL BUFFER
	while True:
		flag, frame = vs.read()
		print(current_frame)
		if flag:
			frame = cv2.resize(frame,(720,480))
			video_frame.append(frame)
			cv2.putText(frame, "Frame: {}".format(current_frame), (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
			cv2.imshow("Frame", frame)
			key = cv2.waitKey(40)

			# pause
			if key == buttons["pause"]:
				print("paused!")
				frame_counter = len(video_frame) - 1
				frame_pause = video_frame[int(frame_counter)]
				while True:
					key2 = cv2.waitKey(0)
					if key2 == buttons["reverse"]:
						if frame_counter > 0:
							frame_counter -= 1
						frame_pause = video_frame[int(frame_counter)]
						cv2.putText(frame_pause, "Paused".format(frame_counter), (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
						cv2.imshow("Frame", frame_pause)
					elif key2 == buttons["pause"]:
						break
					else:
						if frame_counter < len(video_frame) - 1:
							frame_counter += 1
						frame_pause = video_frame[int(frame_counter)]
						cv2.putText(frame_pause, "Paused".format(frame_counter), (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
						cv2.imshow("Frame", frame_pause)
			elif key == buttons["quit"] :
				current_frame += 1
				break

			current_frame += 1
		else:
			vs.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
			print("frame is not ready")
			# It is better to wait for a while for the next frame to be ready
			time.sleep(1.0)

		if current_frame >= 1*int(FPS):
			# If the number of captured frames is equal to the total number of frames,
			# we stop
			break

	while True:
		flag, frame = vs.read()
		if flag:
			frame = cv2.resize(frame,(720,480))
			video_frame.append(frame)
			video_frame.pop(0)
			cv2.putText(frame, "Frame: {}".format(current_frame), (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
			cv2.imshow("Frame", frame)
			key = cv2.waitKey(40)

			# pause
			if key == buttons["pause"]:
				frame_counter = len(video_frame) - 1
				frame_pause = video_frame[int(frame_counter)]
				while True:
					key2 = cv2.waitKey(0)
					if key2 == buttons["reverse"]:
						if frame_counter > 0:
							frame_counter -= 1
						frame_pause = video_frame[int(frame_counter)]
						cv2.putText(frame_pause, "Paused", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
						cv2.imshow("Frame", frame_pause)
					elif key2 == buttons["pause"]:
						break
					else:
						# any other goes forward a frame
						if frame_counter < len(video_frame) - 1:
							frame_counter += 1
						frame_pause = video_frame[int(frame_counter)]
						cv2.putText(frame_pause, "Paused", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
						cv2.imshow("Frame", frame_pause)
			elif key == buttons["quit"] :
				break

			current_frame += 1
		else:
			vs.set(cv2.CAP_PROP_POS_FRAMES, current_frame-1)
			print("frame is not ready")
			# It is better to wait for a while for the next frame to be ready
			time.sleep(1.0)

		if current_frame == num_frames:
			# If the number of captured frames is equal to the total number of frames,
			# we stop
			break

	cv2.destroyAllWindows()
	vs.release()

# arguments: two numpy arrays
# returns a3 such that if a1[row] has nan then a1[row] = a2[row] if a2[row] is not nan
def sync_arrays(a1,a2,SYNC_FRAME):

	t1 = a1[:,0]
	t2 = a2[:,0]
	assert SYNC_FRAME in t1 or SYNC_FRAME in t2

	# remove everything before sync frame
	index1 = list(t1).index(SYNC_FRAME)
	index2 = list(t2).index(SYNC_FRAME)
	a1 = a1[index1:len(a1)]
	a2 = a2[index2:len(a2)]

	# correct size
	if(len(a2) < len(a1)):
		a2 = np.vstack((a2,a1[len(a2):len(a1)]))

	# see returns
	x = np.isnan(a1)
	x = np.where(x)
	x = np.stack(x,axis=1)
	x = x[:,0]
	a1[x] = a2[x]

	return a1

basePath = os.path.dirname(os.path.realpath(__file__)) + "\\"

vidPath = basePath + "vids\\"
csvPath = basePath + "logs\\"
picPath = basePath + "pics\\"
preTrainedPath = basePath + "pre_trained_models\\"

# # get files
# files = []
# for filename in os.listdir(vidPath):
	# print("loaded: " + filename)
	# if(filename.endswith(".MP4") or filename.endswith(".mp4") or filename.endswith(".avi")):
		# files.append(filename)

# lets do some testing
# a1 = np.asarray([[1,1.0],
# [1,1.0],
# [1,1.0],
# [4,nan],
# [5,1.0],
# [6,1.0],
# [7,1.0],
# [8,nan]])

# a2 = np.asarray([[1,1.0],
# [1,5.0],
# [4,5.0],
# [4,5.0],
# [5,5.0],
# [6,5.0],
# [7,5.0],
# [8,5.0]])

# print(a1)
# out = sync_arrays(a1,a2,1)
# print(out)

# test on csvs
txt1 = csvPath + "left_test_ears.csv"
left = pd.read_csv(txt1,sep=',',header=None)

txt2 = csvPath + "right_test_ears.csv"
right = pd.read_csv(txt2,sep=',',header=None)

left = left.values
right = right.values

gc.collect()

out = sync_arrays(left,right,1)
print(out)

# sync videos
# find_sync(files[0])


# coding: utf-8

# In[107]:

import netCDF4 as nc
import numpy as np
import pandas as pd
import sys
import json
import time
import matplotlib.pyplot as plt
from urllib.request import urlopen
from math import sin, cos, sqrt, atan2, radians

# In[26]:

birdfeeder_df = pd.read_csv('PFW2011-12_Subset.csv')

# In[27]:

birdfeeder_df.tail()

# In[50]:

newYork = birdfeeder_df.query('StatProv=="NY"').copy()

# In[51]:

newYork.shape

# In[30]:

ohio = birdfeeder_df.query('StatProv=="OH"')

# In[31]:

ohio.shape

# In[32]:

mass = birdfeeder_df.query('StatProv=="MA"')

# In[33]:

mass.shape

# In[34]:

def calcDistance(lat1, lon1, lat2, lon2):
    R = 6373.0

    latitude1 = radians(lat1)
    longitude1 = radians(lon1)
    latitude2 = radians(lat2)
    longitude2 = radians(lon2)

    dlon = longitude2 - longitude1
    dlat = latitude2 - latitude1

    a = sin(dlat / 2)**2 + cos(latitude1) * cos(latitude2) * sin(dlon / 2)**2
    c = 2 * atan2(sqrt(a), sqrt(1 - a))

    return (R * c)

# In[65]:

newYork = birdfeeder_df.query('StatProv=="NY"').copy()
for index, row in newYork.iterrows():
    dis = calcDistance(row['LATITUDE'], row['LONGITUDE'], 40.730610, -73.935242)
    if (dis > 10):
        newYork.drop(index, inplace=True)

newYork.shape

# In[58]:

newYork.tail()

# In[63]:

boston = birdfeeder_df.query('StatProv=="MA"').copy()
for index, row in boston.iterrows():
    dis = calcDistance(row['LATITUDE'], row['LONGITUDE'], 42.361145, 71.057083)
    if (dis > 10):
        boston.drop(index, inplace=True)

boston.shape

# In[62]:

columbusOH = birdfeeder_df.query('StatProv=="OH"').copy()
for index, row in columbusOH.iterrows():
    dis = calcDistance(row['LATITUDE'], row['LONGITUDE'], 39.949413, -82.959656)
    if (dis > 10):
        columbusOH.drop(index, inplace=True)

columbusOH.shape

# In[64]:

sfCA = birdfeeder_df.query('StatProv=="CA"').copy()
for index, row in sfCA.iterrows():
    dis = calcDistance(row['LATITUDE'], row['LONGITUDE'], 37.733795, -122.446747)
    if (dis > 10):
        sfCA.drop(index, inplace=True)

sfCA.shape

# In[68]:

laCA = birdfeeder_df.query('StatProv=="CA"').copy()
for index, row in laCA.iterrows():
    dis = calcDistance(row['LATITUDE'], row['LONGITUDE'], 34.052235, -118.243683)
    if (dis > 10):
        laCA.drop(index, inplace=True)
laCA.shape

# In[93]:

counter = 0
str1 = 'http://api.wunderground.com/api/f12eaf4e009a36bb/history_'
str2 = '/q/CA/Los_Angeles.json'
laCA['temp'] = np.nan

for index,row in laCA.iterrows():
    if (counter < 10):
        date = row['Year']
        date *= 100
        date += row['Month']
        date *= 100
        date += row['Day']
        date_str = str(date)

        api_link = str1 + date_str + str2
        f = urlopen(api_link)
        json_string = f.read()
        parsed_json = json.loads(json_string)

        if ('history' in parsed_json):
            if ('dailysummary' in parsed_json['history']):
                if (len(parsed_json['history']['dailysummary']) == 1):
                    if ('meantempi' in parsed_json['history']['dailysummary'][0]):
                        temp = parsed_json['history']['dailysummary'][0]['meantempi']
                        laCA.set_value(index, 'temp', temp)

        #temps.append(temp)

        counter += 1
    else:
        time.sleep(60)
        counter = 0

# In[95]:

laCA.shape

# In[96]:

for index, row in laCA.iterrows():
    if(row['temp'] == np.nan):
        laCA.drop(index, inplace=True)

# In[97]:

laCA.shape

# In[99]:

laCA.keys()

# In[100]:

laCA.head()

# In[113]:

laCA.plot(x='temp', y='NSeen', kind='scatter')

# In[114]:

plt.show()

# In[ ]:

counter = 0
str1 = 'http://api.wunderground.com/api/f12eaf4e009a36bb/history_'
str2 = '/q/NY/New_York.json'
newYork['temp'] = np.nan

for index,row in newYork.iterrows():
    if (counter < 10):
        date = row['Year']
        date *= 100
        date += row['Month']
        date *= 100
        date += row['Day']
        date_str = str(date)

        api_link = str1 + date_str + str2
        f = urlopen(api_link)
        json_string = f.read()
        parsed_json = json.loads(json_string)

        if ('history' in parsed_json):
            if ('dailysummary' in parsed_json['history']):
                if (len(parsed_json['history']['dailysummary']) == 1):
                    if ('meantempi' in parsed_json['history']['dailysummary'][0]):
                        temp = parsed_json['history']['dailysummary'][0]['meantempi']
                        newYork.set_value(index, 'temp', temp)
        counter += 1
    else:
        time.sleep(60)
        counter = 0