# $Id: //Depot/Source/SFW/riptable/Python/core/riptable/tests/test_struct.py#14 $


import keyword
import re
import sys
import unittest
from collections import OrderedDict

# from io import StringIO
import numpy as np

import riptable as rt
from riptable import Categorical, Dataset, FastArray, Struct, arange
from riptable.rt_display import DisplayString
from riptable.rt_enum import TypeRegister
from riptable.rt_numpy import arange


class Struct_Test(unittest.TestCase):
    def test_col_ctor_01(self):
        dict1 = {_k: [_i] for _i, _k in enumerate("bac")}
        st = Struct(dict1)
        self.assertEqual(list(st.keys()), list(dict1))
        dict2 = {_k: [_i] for _i, _k in enumerate(["a", "_b", "μεαν"])}
        with self.assertRaises(ValueError):
            st = Struct(dict2)
        with self.assertRaises(TypeError):
            _ = Struct([])
        with self.assertRaises(TypeError):
            _ = Struct(None)
        with self.assertRaises(TypeError):
            _ = Struct(np.array([1]))
        st = Struct()
        self.assertEqual(st.shape, (0, 0))

    # changed invalid name behavior
    # allowing all to pass through, no methods will be lost
    # 04/05/2019
    # def test_col_ctor_02(self):
    #    tempsave = Struct.AllowAnyName
    #    Struct.AllowAnyName=False
    #    hold_WarnOnInvalidNames = Struct.WarnOnInvalidNames
    #    Struct.WarnOnInvalidNames = False
    #    with self.assertRaises(ValueError):
    #        _ = Struct({'a': [0], '0': [1]})
    #    with self.assertRaises(ValueError):
    #        _ = Struct({'a': [0], 'a-b-c': [1]})
    #    with self.assertRaises(ValueError):
    #        # should never be called this way in normal course of events!
    #        struct_instance = Struct({})
    #        struct_instance._validate_names(['a', 'b', 'a'])
    #    # We auto captialize now
    #    #for kwd in keyword.kwlist:
    #    #    with self.assertRaises(ValueError):
    #    #        _ = Struct({'a': [0], kwd: [1]})
    #    #for memb in dir(Struct):
    #    #    with self.assertRaises(ValueError):
    #    #        _ = Struct({'a': [0], memb: [1]})
    #    Struct.WarnOnInvalidNames = hold_WarnOnInvalidNames
    #    Struct.AllowAnyName = tempsave

    def test_col_ctor_02(self):
        inv_keys = ["True", "False", "None"]
        arr = arange(5)
        inv_dict = {k: arr for k in inv_keys}
        with self.assertWarns(UserWarning):
            st = Struct(inv_dict)

        self.assertTrue(bool(np.all(inv_keys == list(st))))
        for k in inv_keys:
            self.assertTrue(bool(np.all(st[k] == arr)))

    def test_col_ctor_03(self):
        tempsave = Struct.AllowAnyName
        Struct.AllowAnyName = False
        hold_WarnOnInvalidNames = Struct.WarnOnInvalidNames
        Struct.WarnOnInvalidNames = True
        with self.assertWarns(UserWarning):
            _ = Struct({"a": [0], "0": [1]})
        with self.assertWarns(UserWarning):
            _ = Struct({"a": [0], "a-b-c": [1]})
        with self.assertRaises(ValueError):  # still should fail even here!
            # should never be called this way in normal course of events!
            struct_instance = Struct({})
            struct_instance._validate_names(["a", "b", "a"])
        for kwd in keyword.kwlist:
            with self.assertWarns(UserWarning):
                _ = Struct({"a": [0], kwd: [1]})
        for memb in dir(Struct):
            with self.assertWarns(UserWarning):
                # We expect really stupid errors resulting from not avoiding invalid names.
                try:
                    _ = Struct({"a": [0], memb: [1]})
                except (TypeError, AttributeError):
                    pass
        Struct.WarnOnInvalidNames = hold_WarnOnInvalidNames
        Struct.AllowAnyName = tempsave

    def test_col_ctor_04(self):
        hold_UseFastArray = Struct.UseFastArray
        arr = np.ones(1000).reshape(5, -1).T
        self.assertFalse(arr.flags.c_contiguous)
        Struct.set_fast_array(True)
        st1 = Struct({"a": arr})
        self.assertFalse(st1.a.flags.c_contiguous)
        Struct.set_fast_array(False)
        st2 = Struct({"a": arr})
        self.assertFalse(st2.a.flags.c_contiguous)
        Struct.set_fast_array(hold_UseFastArray)

    def test_lock(self):
        st1 = Struct({_k: [_i] for _i, _k in enumerate(["a", "b", "c"])})
        st1._lock()
        with self.assertRaises(AttributeError):
            st1.a = 1
        with self.assertRaises(AttributeError):
            st1.d = 1
        with self.assertRaises(AttributeError):
            st1["a"] = 1
        with self.assertRaises(AttributeError):
            st1["d"] = 1
        with self.assertRaises(AttributeError):
            del st1.b
        with self.assertRaises(AttributeError):
            st1.col_remove("c")
        with self.assertRaises(AttributeError):
            _ = st1.col_pop("c")
        with self.assertRaises(AttributeError):
            _ = st1.col_rename("c", "C")
        with self.assertRaises(AttributeError):
            _ = st1.col_map({})
        with self.assertRaises(AttributeError):
            _ = st1.col_move_to_back("c")
        with self.assertRaises(AttributeError):
            _ = st1.col_move_to_front("c")
        st1._unlock()
        st1.a = 1
        st1.d = 1
        del st1.b
        self.assertEqual(list(st1.keys()), ["a", "c", "d"])
        self.assertEqual(st1.a, 1)
        self.assertEqual(st1.d, 1)

    def test_context_interface(self):
        with self.assertRaises(AttributeError):
            with Struct({_k: [_i] for _i, _k in enumerate(["a", "b", "c"])}) as st1:
                self.assertIsInstance(st1, Struct)
                self.assertEqual(list(st1.keys()), ["a", "b", "c"])
                st1._lock()
                st1.a = "cannot modify when locked"
        self.assertEqual(st1.get_ncols(), 0)
        self.assertFalse(st1.is_locked())

    def test_basic_interface(self):
        cols = ["a", "b", "c", "μεαν"]
        dict1 = {_k: [_i] for _i, _k in enumerate(cols)}
        st1 = Struct(dict1)
        self.assertEqual(list(st1.keys()), cols)
        self.assertEqual(list(st1.keys()), cols)
        self.assertEqual(list(st1), cols)
        for _idx, (_k, _v) in enumerate(st1.items()):
            self.assertEqual(st1[_k], dict1[_k])
            self.assertEqual(cols[_idx], _k)
            self.assertEqual(st1[_idx], dict1[_k])
            self.assertEqual(getattr(st1, _k), dict1[_k])
        self.assertEqual(list(st1.keys()), cols)
        self.assertEqual(list(st1), cols)
        self.assertEqual(list(reversed(st1)), list(reversed(cols)))
        self.assertEqual(st1.get_nrows(), 0)
        self.assertEqual(st1.shape, (0, st1.get_ncols()))
        self.assertEqual(len(st1), st1.get_ncols())
        st1["a"] = -999
        self.assertEqual(st1.a, -999)
        st1["newcol"] = -999
        self.assertEqual(st1.newcol, -999)

    def test_col_indexing(self):
        cols = ["b", "a", "c", "μεαν"]
        dict1 = {_k: [_i] for _i, _k in enumerate(cols)}
        st = Struct(dict1)
        self.assertEqual(list(st.keys()), cols)
        self.assertEqual(st["b"], [0])
        self.assertEqual(st[1], [1])
        self.assertEqual(list(st[1:3]), cols[1:3])
        self.assertEqual(list(st[["b", "c"]]), ["b", "c"])
        self.assertEqual(list(st[[True, False, True, False]]), ["b", "c"])
        with self.assertRaises(IndexError):
            _ = st[[True, False, True]]
        with self.assertRaises(IndexError):
            _ = st[[True, False, True, True, True]]
        self.assertEqual(st.b, [0])
        self.assertEqual(st.a, [1])
        self.assertEqual(st.c, [2])
        self.assertEqual(st.μεαν, [3])
        # check shallow copy behavior
        st.d = "stable"
        st1 = st[["a", "b", "d"]]
        st1.a[0] = "diff1"
        st1.b = "diff2"
        st1.d = "diff3"
        self.assertEqual(st1.tolist(), [["diff1"], "diff2", "diff3"])
        self.assertEqual(st.tolist(), [[0], ["diff1"], [2], [3], "stable"])
        with self.assertRaises(AttributeError):
            _ = st.Q
        with self.assertRaises(IndexError):
            _ = st["Q"]
        with self.assertRaises(IndexError):
            _ = st[st.get_ncols()]
        with self.assertRaises(IndexError):
            _ = st[0, "a"]
        with self.assertRaises(IndexError):
            _ = st[["a", "a"]]
        with self.assertRaises(TypeError):
            st[[complex(1)]] = "complex is naughty"
        st1["a"] = "now can do this"
        self.assertEqual(st1.a, "now can do this")
        st1[["a", "b"]] = ["now can do this, too", "as well as this"]
        self.assertEqual(st1.a, "now can do this, too")
        self.assertEqual(st1.b, "as well as this")
        st1[["a", "b"]] = ("a", "b")
        self.assertEqual(st1.a, "a")
        self.assertEqual(st1.b, "b")
        st1["now_there"] = "and even this"
        self.assertEqual(st1.now_there, "and even this")
        with self.assertRaises(IndexError):
            st1[["a", "b"]] = "still cannot do this"
        with self.assertRaises(IndexError):
            st1[["a", "b"]] = np.array(["or", "this"])
        with self.assertRaises(IndexError):
            st1[["a", "b"]] = "AB"

    def test_col_rename(self):
        tempsave = Struct.AllowAnyName
        Struct.AllowAnyName = False
        dict1 = {_k: [_i] for _i, _k in enumerate("abc")}
        st = Struct(dict1)
        with self.assertRaises(ValueError):
            st.col_rename("b", "no-longer-b")
        orig_cols = list(st.keys())
        st.col_rename("a", "a")  # no-op
        with self.assertRaises(ValueError):
            st.col_rename("d", "c")  # cannot rename non-existent column
        self.assertEqual(list(st.keys()), orig_cols)
        with self.assertRaises(ValueError):
            st.col_rename("a", "c")
        for kwd in keyword.kwlist:
            with self.assertRaises(ValueError):
                st.col_rename("a", kwd)
        for memb in dir(Struct):
            with self.assertRaises(ValueError):
                st.col_rename("a", memb)
        st.col_rename("b", "no_longer_b")
        self.assertEqual(list(st.keys()), ["a", "no_longer_b", "c"])
        with self.assertRaises(TypeError):
            st.col_rename("a")
        Struct.AllowAnyName = tempsave

    def test_col_map_01(self):
        orig_cols = ["a", "b", "c"]
        dict1 = {_k: [_i] for _i, _k in enumerate(orig_cols)}
        st = Struct(dict1)
        with self.assertRaises(ValueError):
            st.col_map({"e": "f"})  # cannot rename non-existent column
        self.assertEqual(list(st.keys()), orig_cols)
        st.col_map({_k: "{}_{}".format(_k, _i) for _i, _k in enumerate(st)})
        dict2 = {"{}_{}".format(_k, _i): [_i] for _i, _k in enumerate(orig_cols)}
        st2 = Struct(dict2)
        self.assertEqual(list(st.keys()), list(st2.keys()))
        # dicts no longer equal
        # for _k in st:
        #    self.assertEqual(st.__dict__[_k], st2.__dict__[_k])
        st = Struct(dict1)
        self.assertEqual(list(st.keys()), orig_cols)
        with self.assertRaises(ValueError):
            st.col_map({"a": "e", "b": "d", "c": "d"})
            st.col_map({"q": "z"})  # cannot rename non-existent column
            st.col_map({"q": "a"})  # cannot rename non-existent column
        with self.assertRaises(TypeError):
            st.col_map({"a": "e", "b": "d", "c": 1})
        with self.assertRaises(TypeError):
            st.col_map({1: "e", "b": "d", "c": 1})
        with self.assertRaises(TypeError):
            st.col_map("a")
        self.assertEqual(list(st.keys()), orig_cols)
        st.col_map({"a": "f", "b": "e", "c": "d"})
        self.assertEqual(list(st.keys()), ["f", "e", "d"])
        self.assertEqual(st.f, [0])
        self.assertEqual(st.e, [1])
        self.assertEqual(st.d, [2])

    def test_col_map_02(self):
        st = Struct({_k: list(range(_i * 10, (_i + 1) * 10)) for _i, _k in enumerate("abcdefghijklmnopqrst")})
        cmap = {
            "a": "x",  # 1-transition
            "b": "c",
            "c": "b",  # 2-cycle
            "f": "f",  # 1-cycle
            "g": "h",
            "h": "j",
            "j": "k",
            "k": "g",  # 4-cycle
            "m": "n",
            "n": "o",
            "o": "y",  # 3-transition
            "p": "q",
            "q": "z",
            "z": "r",
        }  # false 3-cycle, really a 2-transition
        st.col_map(cmap)
        for _k in "abcdefghijklmnopqrst":
            self.assertEqual(st[cmap.get(_k, _k)][1], (ord(_k) - ord("a")) * 10 + 1)

    def test_col_map_03(self):
        # same as test_col_map_02 but w/ transitions in scrambled order
        st = Struct({_k: list(range(_i * 10, (_i + 1) * 10)) for _i, _k in enumerate("abcdefghijklmnopqrst")})
        cmap = {
            "a": "x",  # 1-transition
            "b": "c",
            "c": "b",  # 2-cycle
            "f": "f",  # 1-cycle
            "g": "h",
            "k": "g",
            "h": "j",
            "j": "k",  # 4-cycle
            "n": "o",
            "m": "n",
            "o": "y",  # 3-transition
            "q": "z",
            "p": "q",
            "z": "r",
        }  # false 3-cycle, really a 2-transition
        st.col_map(cmap)
        for _k in "abcdefghijklmnopqrst":
            self.assertEqual(st[cmap.get(_k, _k)][1], (ord(_k) - ord("a")) * 10 + 1)

    def test_col_swap(self):
        orig_cols = ["a", "b", "c", "d"]
        dict1 = {_k: [_i] for _i, _k in enumerate(orig_cols)}
        st = Struct(dict1)
        for _i, _k in enumerate(orig_cols):
            self.assertEqual(st[_k][0], _i)
        st.a, st.b = st.b, st.a
        self.assertEqual(st["a"][0], 1)
        self.assertEqual(st["b"][0], 0)
        st.col_swap(["a", "b"], ["b", "a"])
        for _i, _k in enumerate(orig_cols):
            self.assertEqual(st[_k][0], _i)
        st.col_swap(["a", "b", "c"], ["c", "b", "a"])
        self.assertEqual(st["a"][0], 2)
        self.assertEqual(st["b"][0], 1)
        self.assertEqual(st["c"][0], 0)
        self.assertEqual(list(st.keys()), orig_cols)
        st = Struct(dict1)
        st.col_swap(["a", "b", "c", "d"], ["c", "b", "a", "d"])
        self.assertEqual(st["a"][0], 2)
        self.assertEqual(st["b"][0], 1)
        self.assertEqual(st["c"][0], 0)
        self.assertEqual(st["d"][0], 3)
        self.assertEqual(list(st.keys()), orig_cols)

    def test_col_moves(self):
        st = Struct({_k: list(range(_i * 10, (_i + 1) * 10)) for _i, _k in enumerate("abcdefghijklmnop")})

        st.col_move_to_front(1)
        self.assertEqual(list(st), list("bacdefghijklmnop"))
        st.col_move_to_front(1)

        st.col_move_to_back(14)
        self.assertEqual(list(st), list("abcdefghijklmnpo"))
        st.col_move_to_back(14)

        with self.assertRaises(ValueError):
            st.col_move_to_front(arange(20))

        st.col_move_to_back(list("dgh"))
        self.assertEqual(list(st), list("abcefijklmnopdgh"))
        st.col_move_to_front(list("gpha"))
        self.assertEqual(list(st), list("gphabcefijklmnod"))
        st.col_move(list("cim"), list("hfo"))
        self.assertEqual(list(st), list("cimgpabejklndhfo"))
        st.col_move_to_front({"g": 1})
        st.col_move_to_front("h")
        with self.assertWarns(UserWarning):
            st.col_move_to_front("q")
        self.assertEqual(list(st), list("hgcimpabejklndfo"))
        st.col_move_to_back({"g": 1})
        st.col_move_to_back("h")
        with self.assertWarns(UserWarning):
            st.col_move_to_back("q")
        self.assertEqual(list(st), list("cimpabejklndfogh"))

    def test_col_remove(self):
        cols = ["aa", "b", "c", "μεαν"]
        dict1 = {_k: [_i] for _i, _k in enumerate(cols)}
        st1 = Struct(dict1)
        self.assertEqual(list(st1.keys()), cols)
        st1.col_remove("aa")
        self.assertEqual(list(st1.keys()), cols[1:])
        with self.assertRaises(IndexError):
            st1.col_remove("aa")
        st1.col_remove(["b", "c"])
        self.assertEqual(list(st1.keys()), cols[3:])
        with self.assertRaises(IndexError):
            st1.col_remove(["μεαν", "b", "c"])
        with self.assertRaises(TypeError):
            st1.col_remove(2)
        self.assertEqual(list(st1.keys()), cols[3:])
        st1.col_remove(["μεαν"])
        self.assertEqual(list(st1.keys()), [])
        st1.col_remove([])
        self.assertEqual(list(st1.keys()), [])

    def test_col_filter(self):
        x = arange(3)
        ds = Dataset({"a": x, "b": x, "ca": x, "cb": x, "d": x, "e": x, "ear": x})
        self.assertEqual(ds.col_filter(["a", "b"], "c", "e.r").keys(), ["a", "b", "ca", "cb", "ear"])

    def test_col_pop(self):
        cols = ["aa", "b", "c", "μεαν"]
        dict1 = {_k: [_i] for _i, _k in enumerate(cols)}
        st1 = Struct(dict1)
        val = st1.col_pop("aa")
        self.assertEqual(val, [0])
        self.assertEqual(list(st1.keys()), cols[1:])
        with self.assertRaises(IndexError):
            st1.col_pop("aa")
        with self.assertRaises(IndexError):
            st1.col_pop(["aa"])
        st2 = st1.col_pop(["b", "c"])
        self.assertEqual(list(st1.keys()), cols[3:])
        self.assertEqual(list(st2.keys()), ["b", "c"])
        self.assertEqual(st2.b, [1])
        self.assertEqual(st2.c, [2])
        st3 = st1.col_pop(slice(None))
        self.assertEqual(st1.get_ncols(), 0)
        self.assertEqual(st3.get_ncols(), 1)
        self.assertEqual(st3.μεαν, [3])

    def test_col_str_match(self):
        cols = ["apple", "orange", "kumquat", "koala", "μεαν", "two_kumquats_or_more"]
        dict1 = {_k: [_i] for _i, _k in enumerate(cols)}
        st1 = Struct(dict1)
        self.assertEqual(st1.col_str_match(r".*kum.*t").shape, (st1.get_ncols(),))
        self.assertEqual(st1.col_str_match(r".*kum.*t").sum(), 2)
        self.assertEqual(
            list(st1[st1.col_str_match(r".*kum.*t")]),
            ["kumquat", "two_kumquats_or_more"],
        )
        self.assertEqual(
            list(st1[st1.col_str_match(r".*KuM.*t", flags=re.IGNORECASE)]),
            ["kumquat", "two_kumquats_or_more"],
        )

    @staticmethod
    def _sum_dict(dd):
        return {_k: sum(_v) for _k, _v in dd.items()}

    def test_as_dictionary(self):
        od = {_k: list(range(_i * 10, (_i + 1) * 10)) for _i, _k in enumerate("abcdefghij")}
        st = Struct(od)
        dt1 = st.asdict()
        dt2 = st.asdict(copy=True)
        self.assertEqual(self._sum_dict(dt1), self._sum_dict(od))
        self.assertIsInstance(dt1, dict)
        self.assertEqual(list(st.asdict(sublist=["a", "c", "d"]).keys()), ["a", "c", "d"])
        with self.assertRaises(AttributeError):
            _ = st.asdict(sublist=["a", "c", "q"]).keys()
        # test for copy/no-copy?
        lst10 = list(range(10))
        self.assertEqual(st.a, lst10)
        self.assertEqual(dt1["a"], lst10)
        self.assertEqual(dt2["a"], lst10)
        st.a[5] = -1
        self.assertEqual(st.a[5], -1)
        self.assertEqual(dt1["a"][5], -1)
        self.assertEqual(dt2["a"], lst10)
        dt2["b"][3] = -2
        self.assertEqual(st.b[3], 13)
        self.assertEqual(dt1["b"][3], 13)

    def test_as_ordered_dictionary(self):
        od = {_k: list(range(_i * 10, (_i + 1) * 10)) for _i, _k in enumerate("abcdefghij")}
        st = Struct(od)
        dt1 = st.as_ordered_dictionary()
        self.assertEqual(self._sum_dict(dt1), self._sum_dict(od))
        self.assertIsInstance(dt1, OrderedDict)
        self.assertEqual(
            list(st.as_ordered_dictionary(sublist=["a", "c", "d"]).keys()),
            ["a", "c", "d"],
        )
        with self.assertRaises(ValueError):
            _ = st.as_ordered_dictionary(sublist=["a", "c", "q"]).keys()

    def test_no_bool(self):
        with self.assertRaises(ValueError):
            _ = bool(Struct({"A": 5}))
        with self.assertRaises(ValueError):
            _ = bool(Struct({"A": [5.5], "B": "fish"}))
        with self.assertRaises(ValueError):
            _ = bool(Struct({"A": True}))
        with self.assertRaises(ValueError):
            _ = bool(Struct({"A": True, "B": False}))

    def test_tolist(self):
        data = {"a": 5, "b": 5.6, "c": "fish", "d": "μεαν", "e": True, "f": {"A": None}}
        st = Struct(data)
        self.assertEqual(st.tolist(), list(data.values()))

    def test_any(self):
        self.assertTrue(Struct({"A": True, "B": True}).any())
        self.assertTrue(Struct({"A": True, "B": False}).any())
        self.assertFalse(Struct({"A": False, "B": False}).any())
        self.assertTrue(Struct({"A": True}).any())
        self.assertFalse(Struct({"A": False}).any())
        self.assertTrue(Struct({"A": [5.5], "B": "fish"}).any())
        self.assertTrue(Struct({"A": np.nan, "B": False}).any())
        self.assertTrue(Struct({"A": ["A", "B", ""], "B": False}).any())
        self.assertTrue(Struct({"A": np.array(["A", "B", ""]), "B": False}).any())
        self.assertFalse(Struct({}).any())  # there does not exist one which is true

    def test_all(self):
        self.assertTrue(Struct({"A": True, "B": True}).all())
        self.assertFalse(Struct({"A": True, "B": False}).all())
        self.assertFalse(Struct({"A": False, "B": False}).all())
        self.assertTrue(Struct({"A": True}).all())
        self.assertFalse(Struct({"A": False}).all())
        self.assertTrue(Struct({"A": [5.5], "B": "fish"}).all())
        self.assertFalse(Struct({"A": np.nan, "B": False}).all())
        self.assertFalse(Struct({"A": ["A", "B", ""], "B": False}).all())
        self.assertFalse(Struct({"A": np.array(["A", "B", ""]), "B": False}).all())
        self.assertTrue(Struct({}).all())  # all that exist are true

    def test_comparison_01(self):
        st1 = Struct({"A": [5.5], "B": "fish"})
        st2 = Struct({"A": [5.5], "B": "fish", "C": [380, 20]})
        st3 = Struct({"B": "fish", "C": [380, 20]})
        st4 = Struct({"A": [5.5], "B": np.ones(4), "C": [380, 20]})
        self.assertEqual((st1 == st1).asdict(), {"A": True, "B": True})
        self.assertEqual((st1 == st2).asdict(), {"A": True, "B": True, "C": False})
        self.assertEqual((st2 == st1).asdict(), {"A": True, "B": True, "C": False})
        self.assertEqual(list((st2 == st3).keys()), ["A", "B", "C"])
        self.assertEqual(list((st3 == st2).keys()), ["B", "C", "A"])
        self.assertEqual((st4 == st4).asdict(), {"A": True, "B": True, "C": True})
        for comp in "eq ne lt le gt ge".split():
            self.assertIsInstance((getattr(st1, f"__{comp}__")(st2)).all(), bool)

    def test_comparison_02(self):
        st1 = Struct({"A": [5.5], "B": 20})
        st2 = Struct({"A": [5.5], "B": "20"})
        st3 = Struct({"A": [6.5], "B": 30})
        self.assertEqual((st1 < st3).asdict(), {"A": True, "B": True})
        with self.assertRaises(NotImplementedError):
            _ = st2 < st3
        with self.assertRaises(TypeError):
            _ = st2 < []

    def test_concat(self):
        st1 = Struct(
            {
                "ds": TypeRegister.Dataset({"col_" + str(i): np.random.rand(5) for i in range(5)}),
                "arr": arange(5),
                "cat": TypeRegister.Categorical(["a", "a", "b", "c", "a"]),
            }
        )
        st2 = Struct(
            {
                "ds": TypeRegister.Dataset({"col_" + str(i): np.random.rand(5) for i in range(5)}),
                "arr": arange(5),
                "cat": TypeRegister.Categorical(["a", "a", "b", "c", "a"]),
            }
        )

        result = Struct.concat_structs([st1, st2])
        self.assertTrue(isinstance(result.ds, TypeRegister.Dataset))
        self.assertTrue(isinstance(result.arr, TypeRegister.FastArray))
        self.assertTrue(isinstance(result.cat, TypeRegister.Categorical))

        correct_arr = np.hstack([st1.arr, st2.arr])
        self.assertTrue(bool(np.all(correct_arr == result.arr)))

        correct_cat = np.array(["a", "a", "b", "c", "a", "a", "a", "b", "c", "a"])
        self.assertTrue(bool(np.all(correct_cat == result.cat)))

        for c in st1.ds:
            correct = np.hstack([st1.ds[c], st2.ds[c]])
            self.assertTrue(bool(np.all(correct == result.ds[c])))

        order = list(st1.keys())
        result = list(result.keys())
        self.assertTrue(bool(np.all(order == result)))

    def test_categorical(self):
        ds = Dataset({"A": ["a", "b"], "B": ["c", "d"], "C": ["a", "μεαν"]})
        # to_categorical called the categorical constructor
        # removed, because the user can simply say c=Categorical(ds.A)
        # self.assertIsInstance(Dataset.to_categorical(ds.A), Categorical)
        # self.assertNotIsInstance(ds.A, Categorical)
        ds.make_categoricals("A")
        self.assertIsInstance(ds.A, Categorical)
        ds = Dataset({"A": ["a", "b"], "B": ["c", "d"], "C": ["a", "μεαν"]})
        # 9/28/2018 SJK: dataset no longer flips unicode to categorical
        # self.assertIsInstance(ds.C, Categorical)
        ds.make_categoricals(["A", "C"])
        self.assertIsInstance(ds.A, Categorical)
        self.assertIsInstance(ds.C, Categorical)
        # BUG: categorical is not retaining its base index for true unicode...
        self.assertEqual(ds.C[1], "μεαν")
        ds = Dataset({"A": ["a", "b"], "B": ["c", "d"], "C": ["a", "μεαν"], "D": [1, 2]})
        types1 = {"A": FastArray, "B": FastArray, "C": FastArray, "D": FastArray}
        for _k in ds:
            self.assertIsInstance(ds[_k], types1[_k])
        ds.make_categoricals()
        types2 = {"A": Categorical, "B": Categorical, "C": Categorical, "D": FastArray}
        for _k in ds:
            self.assertIsInstance(ds[_k], types2[_k])

        ds1 = Dataset(
            {
                "a": np.random.choice(["a", "b", "c"], 5),
                "b": np.random.choice(["a", "b", "c"], 5),
            }
        )
        ds2 = Dataset(
            {
                "a": np.random.choice(["a", "b", "c"], 5),
                "b": np.random.choice(["a", "b", "c"], 5),
            }
        )
        ds3 = Dataset(
            {
                "a": np.random.choice(["a", "b", "c"], 5),
                "b": np.random.choice(["a", "b", "c"], 5),
            }
        )
        arr = np.random.choice(["a", "b", "c"], 5)
        st = Struct({"ds1": ds1, "ds2": ds2, "ds3": ds3, "arr": arr})
        st.make_categoricals()
        for ds in [st.ds1, st.ds2, st.ds3]:
            for col in ds:
                self.assertTrue(isinstance(ds[col], Categorical))
        self.assertTrue(isinstance(st.arr, Categorical))

    def _test_output(self):
        if sys.platform != "win32":
            intt_name = b"int64"
        else:
            intt_name = b"int32"

        class Dataset1(Struct):
            pass  # dummy for testing, mimics behavior of real Dataset

        st = Struct(
            {
                "a": Dataset1({"A": range(10), "B": range(10, 20)}),
                "b": Struct({"C": 0, "D": 1, "E": 2}),
                "c": FastArray(np.arange(5)),
                "d": np.arange(5, 10),
                "e": ["abc", "def", "ghi"],
                "f": {"q": 1, "r": 2},
                "g": 3.14,
                "h": 84,
                "i": None,
                "j": slice(None),
            }
        )
        headers, spec = st.get_table_data()
        self.assertEqual(len(headers), 1)
        self.assertEqual([hd.col_name for hd in headers[0]], ["Name", "Type", "Rows", "0", "1", "2"])
        self.assertEqual(
            [_r.tolist() for _r in spec],
            [
                [b"a", b"b", b"c", b"d", b"e", b"f", b"g", b"h", b"i", b"j"],
                [
                    b"Dataset1",
                    b"Struct",
                    intt_name,
                    intt_name,
                    b"list",
                    b"dict",
                    b"float",
                    b"int",
                    b"NoneType",
                    b"slice",
                ],
                [b"2", b"3", b"5", b"5", b"3", b"2", b"0", b"0", b"0", b"0"],
                [b"A", b"", b"0", b"5", b"", b"", b"3.14", b"84", b"", b""],
                [b"B", b"", b"1", b"6", b"", b"", b"", b"", b"", b""],
                [b"", b"", b"2", b"7", b"", b"", b"", b"", b"", b""],
            ],
        )
        self.assertEqual(
            str(st),
            f"""#   Name   Type       Rows   0      1   2
-   ----   --------   ----   ----   -   -
0   a      Dataset1   2      A      B
1   b      Struct     3
2   c      {intt_name.decode()}      5      0      1   2
3   d      {intt_name.decode()}      5      5      6   7
4   e      list       3
5   f      dict       2
6   g      float      0      3.14
7   h      int        0      84
8   i      NoneType   0
9   j      slice      0                  """,
        )
        self.assertEqual(Struct._sizeof_fmt(128), "128.0 B")
        tsize = 1280
        for unit in ["K", "M", "G", "T", "P", "E", "Z", "Y"]:
            self.assertEqual(Struct._sizeof_fmt(tsize), f"1.2 {unit}B")
            tsize *= 1024
        self.assertEqual(st._last_row_stats(), "[10 columns]")

    def test_total_sizes(self):
        st = Struct(
            {
                "a": Dataset(
                    {
                        # 10x int32 => 40B
                        "A": range(10),
                        # 10x int32 => 40B
                        "B": range(10, 20),
                    }
                ),
                "b": Struct(
                    {
                        # 1x int32 => 4B
                        "C": 0,
                        # 1x int32 => 4B
                        "D": 1,
                        # 1x int32 => 4B
                        "E": 2,
                    }
                ),
                # 5x int32 => 20B
                "c": FastArray(np.arange(5)),
                # 5x int32 => 20B
                "d": np.arange(5, 10),
                # ???
                "e": ["abc", "def", "ghi"],
                "f": {
                    # 1x int32 => 4B
                    "q": 1,
                    # 1x int32 => 4B
                    "r": 2,
                },
                # 1x float64 => 8B
                "g": 3.14,
                # 1x int32 => 4B
                "h": 84,
                # ???
                "i": None,
                # ???
                "j": slice(None),
            }
        )

        # Create some duplicated/aliased data within the struct.
        st.z = st.c

        # Calculate the sizes of the Struct's data in bytes.
        (physical, logical) = st.total_sizes

        # For now, we only check that the logical size is larger than the physical size
        # (due to the presence of aliased array(s) somewhere within the Struct).
        # TODO: Strengthen this test by checking the actual computed sizes to make sure they're correct.
        self.assertLess(physical, logical, "The physical size is not less than the logical size.")

    def test_total_sizes_with_categorical(self):
        st = Struct({"c": Categorical(["aa", "bb", "cc", "dd"])})
        st.d = st.c
        (physical, logical) = st.total_sizes
        self.assertEqual(physical, logical // 2)
        self.assertGreaterEqual(physical, np.asarray(st.c).nbytes + st.c.category_array.nbytes)

    def test_tree(self):
        """sanity check that .tree() at least returns something, even for empty Struct"""
        s = Struct()
        self.assertIsInstance(s.tree(), DisplayString)
        s["foo"] = Dataset({"bar": arange(5)})
        self.assertIsInstance(s.tree(), DisplayString)

    def test_key_search(self):
        st = Struct()
        st.column1 = rt.arange(10)
        st.Apple = rt.arange(10)
        myds = rt.Dataset()
        myds.animal = rt.arange(10)
        myds.ale = rt.arange(10)
        st.opp = myds
        mystruct = rt.Struct()
        mystruct.allen = rt.arange(10)
        substruct = rt.Struct()
        myds2 = rt.Dataset()
        myds2.oval = rt.arange(10)
        myds2.elena = rt.arange(10)
        myds2.Valence = rt.arange(10)
        substruct.myds2 = myds2
        mystruct.substruct = substruct
        st.Alex = mystruct
        self.assertEqual(
            st.key_search("a.*e"), ["Apple", "Alex", "opp.ale", "Alex.allen", "Alex.substruct.myds2.Valence"]
        )


if __name__ == "__main__":
    tester = unittest.main()
