import json
import os

from tempfile import TemporaryDirectory
from unittest import TestCase

import numpy

from numpy.testing import assert_allclose

from cogent3 import DNA
from cogent3.util.dict_array import (
    DictArray,
    DictArrayTemplate,
    convert2DDict,
    convert2Ddistance,
    convert_1D_dict,
    convert_dict,
    convert_for_dictarray,
    convert_series,
)


class DictArrayTest(TestCase):
    a = numpy.identity(3, int)

    def test_convert_series(self):
        """convert_series produces valid template input"""
        vals, row_keys, col_keys = convert_series([[4], [5]], ["A", "B"], ["a"])
        b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        self.assertEqual(b.array.tolist(), [[4], [5]])
        data = [[245, 599]]
        vals, row_keys, col_keys = convert_series(data)
        b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        self.assertEqual(b.array.tolist(), data)

        vals, row_keys, col_keys = convert_series(data[0])
        b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        self.assertEqual(b.array.tolist(), data[0])

    def test_convert_dict(self):
        """convert_dict produces valid template input"""
        twoDdict = dict(a=dict(b=4, c=5))
        vals, row_keys, col_keys = convert_dict(twoDdict)
        b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        self.assertEqual(b.array.tolist(), [[4, 5]])

    def test_convert2DDict(self):
        """convert2DDict produces valid template input"""
        data = dict(a=dict(b=4, c=5))
        vals, row_keys, col_keys = convert2DDict(data)
        self.assertEqual(set(row_keys), set(["a"]))
        b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        self.assertEqual(b.array.tolist(), [[4, 5]])
        # row keys, then column
        self.assertEqual(b.template.names, [["a"], ["b", "c"]])

        data = {
            "a": {"a": 0, "b": 1, "e": 0},
            "b": {"a": 1, "b": 0, "e": 4},
            "e": {"a": 0, "b": 4, "e": 0},
        }
        vals, row_keys, col_keys = convert2DDict(data)
        b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        got = b.to_dict()
        self.assertEqual(got, data)
        self.assertEqual(b.template.names, [["a", "b", "e"], ["a", "b", "e"]])

        data = dict(a=dict(b=4, c=5))
        vals, row_keys, col_keys = convert2DDict(data, make_symmetric=True)
        self.assertEqual(row_keys, col_keys)
        self.assertEqual(vals, [[0, 4, 5], [4, 0, 0], [5, 0, 0]])

    def test_convert2Ddistance(self):
        """convert2Ddistance produces valid template input"""
        data = {("a", "b"): 4, ("a", "c"): 5}
        vals, row_keys, col_keys = convert2Ddistance(data)
        b = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        self.assertEqual(b.array.tolist(), [[0, 4, 5], [4, 0, 0], [5, 0, 0]])

    def test_convert_1D_dict(self):
        """convert_1D_dict produces valid template input"""
        data = dict(a=0, b=35, c=45)
        vals, keys = convert_1D_dict(data)
        b = DictArrayTemplate(keys)
        b = b.wrap(vals)
        self.assertEqual(b.array.tolist(), [0, 35, 45])

    def test_construct_both_dim_str(self):
        """correctly construct when keys for both dimensions are str"""
        b = DictArrayTemplate("abc", "ABC").wrap(self.a)
        self.assertEqual(b[0].array.tolist(), [1, 0, 0])
        self.assertEqual(b["a"].array.tolist(), [1, 0, 0])
        self.assertEqual(b.array.tolist(), [[1, 0, 0], [0, 1, 0], [0, 0, 1]])

    def test_key_levels(self):
        """DictArray both levels have keys."""
        b = DictArrayTemplate("abc", "ABC").wrap(self.a)
        self.assertEqual(b.keys(), ["a", "b", "c"])
        self.assertEqual(b["a"].keys(), ["A", "B", "C"])
        self.assertEqual(list(b["a"]), [1, 0, 0])
        self.assertEqual(sum(b["a"]), 1)

    def test_int_labels(self):
        """DictArray with no labels."""
        b = DictArrayTemplate(3, 3).wrap(self.a)
        self.assertEqual(b.keys(), [0, 1, 2])
        self.assertEqual(b[0].keys(), [0, 1, 2])
        self.assertEqual(sum(b[0]), 1)

    def test_str_labels(self):
        """DictArray with str labels or numpy U dtype"""
        b = DictArrayTemplate(["Ab", "Bb", "Cb"], ["Db", "Eb", "Fb"]).wrap(self.a)
        self.assertEqual(b.keys(), ["Ab", "Bb", "Cb"])
        self.assertEqual(b[0].keys(), ["Db", "Eb", "Fb"])
        self.assertEqual(b["Ab", "Eb"], 0)
        self.assertEqual(b["Bb", "Eb"], 1)

        b = DictArrayTemplate(
            numpy.array(["Ab", "Bb", "Cb"], dtype="U"),
            numpy.array(["Db", "Eb", "Fb"], dtype="U"),
        ).wrap(self.a)
        self.assertEqual(b.keys(), ["Ab", "Bb", "Cb"])
        self.assertEqual(b[0].keys(), ["Db", "Eb", "Fb"])
        self.assertEqual(b["Ab", "Eb"], 0)
        self.assertEqual(b["Bb", "Eb"], 1)

    def test_with_mixed_label_types(self):
        """DictArray constructed with mixed label types."""
        b = DictArrayTemplate("ABC", 3).wrap(self.a)
        self.assertEqual(b.keys(), ["A", "B", "C"])
        self.assertEqual(b["A"].keys(), [0, 1, 2])

    def test_numpy_ops(self):
        """DictArray should work properly in numpy operations."""
        darr = DictArrayTemplate(list(DNA), list(DNA)).wrap(
            [
                [0.7, 0.1, 0.1, 0.1],
                [0.1, 0.7, 0.1, 0.1],
                [0.1, 0.1, 0.7, 0.1],
                [0.1, 0.1, 0.1, 0.7],
            ]
        )
        mprobs = numpy.array([0.25, 0.25, 0.25, 0.25])
        assert_allclose(mprobs.dot(darr), [0.25, 0.25, 0.25, 0.25])
        assert_allclose(numpy.dot(mprobs, darr), [0.25, 0.25, 0.25, 0.25])

    def test_to_dict(self):
        """DictArray should convert 1D / 2D arrays with/without named row"""
        # 1D data, only 1D keys provided
        data = [0, 35, 45]
        keys = "a", "b", "c"
        darr = DictArrayTemplate(keys).wrap(data)
        self.assertEqual(darr.to_dict(), dict(zip(keys, data)))
        # 2D data, 2D keys, both string, provided
        data = [[0, 35, 45]]
        darr = DictArrayTemplate(["0"], keys).wrap(data)
        darr.to_dict()
        self.assertEqual(darr.to_dict(), {"0": {"a": 0, "b": 35, "c": 45}})
        # 2D data, 2D keys, one int, one string, provided
        darr = DictArrayTemplate([1], keys).wrap(data)
        self.assertEqual(darr.to_dict(), {1: {"a": 0, "b": 35, "c": 45}})
        darr = DictArrayTemplate([0], keys).wrap(data)
        self.assertEqual(darr.to_dict(), {0: {"a": 0, "b": 35, "c": 45}})

    def test_to_dict_1d(self):
        """should successfully produce a 1D dict"""
        data = {
            "ABAYE2984": {
                "ABAYE2984": 0,
                "Atu3667": None,
                "Avin_42730": 0.6381173875591908,
                "BAA10469": None,
            },
            "Atu3667": {
                "ABAYE2984": None,
                "Atu3667": 0,
                "Avin_42730": 2.3682377869318993,
                "BAA10469": None,
            },
            "Avin_42730": {
                "ABAYE2984": 0.6381173875591908,
                "Atu3667": 2.3682377869318993,
                "Avin_42730": 0,
                "BAA10469": 1.8515731266342546,
            },
            "BAA10469": {
                "ABAYE2984": None,
                "Atu3667": None,
                "Avin_42730": 1.8515731266342546,
                "BAA10469": 0,
            },
        }
        darr = DictArray(data, dtype="O")
        expect = {
            (n1, n2): darr[n1, n2]
            for n1 in darr.template.names[0]
            for n2 in darr.template.names[1]
        }
        self.assertEqual(darr.to_dict(flatten=True), expect)

        darr = DictArrayTemplate(["s1", "s2"], ["s1", "s2"]).wrap(
            [[0.0, 0.25], [0.25, 0.0]]
        )
        self.assertEqual(
            darr.to_dict(flatten=True),
            {
                ("s1", "s2"): 0.25,
                ("s2", "s1"): 0.25,
                ("s1", "s1"): 0.0,
                ("s2", "s2"): 0.0,
            },
        )

    def test_to_dict_nested(self):
        """DictArray.to_dict() should convert nested DictArray instances to
        dict's too."""
        a = numpy.identity(3, int)
        b = DictArrayTemplate("abc", "ABC")
        b = b.wrap(a)
        self.assertEqual(b.array.tolist(), [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
        c = DictArrayTemplate("de", "DE").wrap([[b, b], [b, b]])
        self.assertTrue(isinstance(c.to_dict()["d"], dict))

    def test_to_dict_values(self):
        """values from to_dict should be python types"""
        keys = "a", "b", "c", "d"
        for data, _type in [
            ([0, 35, 45, 3], int),
            (["abc", "def", "jkl;", "aa"], str),
            ([0.1, 0.2, 0.3, 0.4], float),
        ]:
            darr = DictArrayTemplate(keys).wrap(data)
            got = {type(v) for v in darr.to_dict().values()}
            self.assertEqual(got, {_type})

        for data, _type in [
            ([0, 35, 45, 3], int),
            (["abc", "def", "jkl;", "aa"], str),
            ([0.1, 0.2, 0.3, 0.4], float),
        ]:
            darr = DictArrayTemplate(keys[:2], keys[2:]).wrap([data[:2], data[2:]])
            got = {type(v) for d in darr.to_dict().values() for v in d.values()}
            self.assertEqual(got, {_type})

    def test_to_dict_json(self):
        """should be able to json.dumps result of to_dict"""
        keys = "a", "b", "c", "d"
        for data in [
            [0, 35, 45, 3],
            ["abc", "def", "jkl;", "aa"],
            [0.1, 0.2, 0.3, 0.4],
        ]:
            darr = DictArrayTemplate(keys).wrap(data)
            got = json.dumps(darr.to_dict())
            self.assertIsInstance(got, str)

        for data in [
            [0, 35, 45, 3],
            ["abc", "def", "jkl;", "aa"],
            [0.1, 0.2, 0.3, 0.4],
        ]:
            darr = DictArrayTemplate(keys[:2], keys[2:]).wrap([data[:2], data[2:]])
            got = json.dumps(darr.to_dict())
            self.assertIsInstance(got, str)

    def test_to_dict_roundtrip(self):
        """roundtrip of DictArray.to_dict() should produce same order."""
        d1 = dict(a=dict(k=1, l=2, m=3), b=dict(k=4, l=5, m=6))
        darr1 = DictArray(d1)
        d2 = darr1.to_dict()
        self.assertEqual(d1, d2)
        d3 = DictArray(d2)
        self.assertEqual(d1, d3)

    def test_convert_for_dictarray(self):
        """successfully delegates when constructed from a DictArray"""
        a = numpy.identity(3, int)
        b = DictArrayTemplate("abc", "ABC").wrap(a)
        vals, row_keys, col_keys = convert_for_dictarray(b)
        got = DictArrayTemplate(row_keys, col_keys).wrap(vals)
        self.assertEqual(got.array.tolist(), b.array.tolist())
        # the wrap method creates a new array
        self.assertIsNot(got.array, b.array)

    def test_convert_for_dictarray2(self):
        """convert_for_dictarray correctly delegates"""
        b = DictArrayTemplate("abc", "ABC").wrap(self.a)
        data_types = (
            [[245, 599]],
            dict(a=dict(b=4, c=5)),
            {("a", "b"): 4, ("a", "c"): 5},
            dict(a=0, b=35, c=45),
            b,
        )
        for data in data_types:
            vals, row_keys, col_keys = convert_for_dictarray(data)
            _ = DictArrayTemplate(row_keys, col_keys).wrap(vals)

    def test_direct_construction(self):
        """directly construct a dict array"""
        b = DictArrayTemplate("abc", "ABC").wrap(self.a)
        data_types = (
            [[245, 599]],
            dict(a=dict(b=4, c=5)),
            {("a", "b"): 4, ("a", "c"): 5},
            dict(a=0, b=35, c=45),
            b,
        )
        # exercising construction
        for data in data_types:
            _ = DictArray(data)

    def test_getitem(self):
        """correctly slices"""
        darr = DictArrayTemplate(list(DNA), list(DNA)).wrap(
            [
                [0.7, 0.1, 0.1, 0.1],
                [0.1, 0.7, 0.1, 0.1],
                [0.1, 0.1, 0.7, 0.1],
                [0.1, 0.1, 0.1, 0.7],
            ]
        )
        r = darr[:, "A":"G"]
        assert_allclose(r.to_array(), [[0.1], [0.1], [0.7], [0.1]])
        r = darr[2:, "A":"G"]
        assert_allclose(r.to_array(), [[0.7], [0.1]])

    def test_to_normalized(self):
        """computes frequencies across correct dimension"""
        data = [[3, 7], [2, 8], [5, 5]]
        darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
        row_normal = darr.to_normalized(by_row=True)
        assert_allclose(row_normal.array, [[0.3, 0.7], [0.2, 0.8], [0.5, 0.5]])
        col_normal = darr.to_normalized(by_column=True)
        assert_allclose(col_normal.array, [[0.3, 7 / 20], [0.2, 8 / 20], [0.5, 5 / 20]])
        # trying to do both raises AssertionError
        with self.assertRaises(AssertionError):
            darr.to_normalized(by_row=True, by_column=True)

    def test_col_sum(self):
        """correctly computes column sum"""
        data = [[3, 7], [2, 8], [5, 5]]
        darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
        col_sum = darr.col_sum()
        assert_allclose(col_sum.array, [10, 20])

    def test_row_sum(self):
        """correctly computes row sum"""
        data = [[3, 7], [2, 8], [5, 5]]
        darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
        row_sum = darr.row_sum()
        assert_allclose(row_sum.array, [10, 10, 10])

    def test_get_repr_html(self):
        """exercising method used by parent classes for nice Jupyter display"""
        data = [[3, 7], [2, 8], [5, 5]]
        darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
        got = darr._repr_html_()
        self.assertIsInstance(got, str)
        self.assertTrue(len(got), 100)

        # case where 1D array
        a = [4, 6, 4, 2]
        darr = DictArrayTemplate(["A", "C", "G", "T"]).wrap(a)
        got = darr._repr_html_()
        self.assertTrue('class="index"' not in got)

        # case of 3D array
        d3 = numpy.arange(8).reshape((2, 2, 2))
        darr = DictArrayTemplate(2, 2, 2).wrap(d3)
        got = darr._repr_html_()
        self.assertIn("3 dimensional", got)

    def test_write(self):
        """exercising write method"""
        data = [[3, 7], [2, 8], [5, 5]]
        darr = DictArrayTemplate(list("ABC"), list("ab")).wrap(data)
        with TemporaryDirectory(dir=".") as dirname:
            outpath = os.path.join(dirname, "delme.tsv")
            darr.write(outpath)
            with open(outpath) as infile:
                contents = [l.strip().split() for l in infile]
            header = contents.pop(0)
            self.assertEqual(header, ["dim-1", "dim-2", "value"])
            got = {(k1, k2): int(v) for k1, k2, v in contents}
            self.assertEqual(got, darr.to_dict(flatten=True))

    def test_to_string(self):
        darr = DictArrayTemplate(2, 2).wrap(
            [[3.123456789, 2 * 3.123456789], [3 * 3.123456789, 4 * 3.123456789]]
        )
        self.assertEqual(
            darr.to_string(sep=","),
            "dim-1,dim-2,value\n0,0,3.123456789\n0,1,6.246913578\n1,0,9.370370367\n1,1,12.493827156",
        )
        self.assertEqual(
            darr.to_string(),
            "dim-1\tdim-2\tvalue\n0\t0\t3.123456789\n0\t1\t6.246913578\n1\t0\t9.370370367\n1\t1\t12.493827156",
        )
        self.assertEqual(
            darr.to_string(sep=" "),
            "dim-1 dim-2 value\n0 0 3.123456789\n0 1 6.246913578\n1 0 9.370370367\n1 1 12.493827156",
        )
        with self.assertRaises(ValueError):
            darr.to_string(format="md"),

    def test_to_table(self):
        """creates Table when ndim <= 2"""
        from cogent3.util.table import Table

        a1D = DictArrayTemplate(["a", "b"]).wrap([0, 1])
        t = a1D.to_table()
        self.assertIsInstance(t, Table)
        # 1D tables don't get an index_name column
        self.assertEqual(t.index_name, None)
        a2D = DictArrayTemplate(["a", "b"], ["c", "d"]).wrap(
            numpy.array([0, 1, 2, 3]).reshape((2, 2))
        )
        t = a2D.to_table()
        self.assertIsInstance(t, Table)
        self.assertEqual(t.shape, (2, 3))  # because index_name column added
        # make sure the 2D variant has an index_name column, name is empty string
        self.assertEqual(t.index_name, "")
        self.assertEqual(t.columns[""].tolist(), a2D.template.names[0])
        # which works
        self.assertEqual(t["b", "d"], 3)

        a3D = DictArrayTemplate(["a", "b"], ["c", "d"], ["e", "f"]).wrap(
            numpy.array([0, 1, 2, 3, 4, 5, 6, 7]).reshape((2, 2, 2))
        )
        with self.assertRaises(ValueError):
            _ = a3D.to_table()

    def test_interpret_index(self):
        """correctly handles just explicitly defined indices"""
        n = ["ab", "dna", "rna"]
        a1D = DictArrayTemplate(n)
        got = a1D.interpret_index(["ab", "rna"])
        self.assertEqual(got[0], ([0, 2],))
        got = a1D.interpret_index([0, 2])
        self.assertEqual(got[0], ([0, 2],))

    def test_slicing_combos(self):
        """different mixtures of slicing should work"""
        darr = DictArrayTemplate(list(DNA), list(DNA)).wrap(
            [
                [0.7, 0.1, 0.2, 0.3],
                [0.1, 0.7, 0.1, 0.3],
                [0.3, 0.2, 0.6, 0.3],
                [0.4, 0.1, 0.1, 0.7],
            ]
        )
        got = darr["C":"G", "C":"G"]
        assert_allclose(got.array, numpy.array([[0.7, 0.1], [0.2, 0.6]]))

        got = darr[[1, 2], [1, 2]]
        assert_allclose(got.array, numpy.array([[0.7, 0.1], [0.2, 0.6]]))

        got = darr[[2, 3], "C"]
        assert_allclose(got.array, numpy.array([0.2, 0.1]))
        got = darr["C", [2, 3]]
        assert_allclose(got.array, numpy.array([0.1, 0.3]))

        got = darr[[1, 2], "T":"A"]
        assert_allclose(got.array, numpy.array([[0.1, 0.7], [0.3, 0.2]]))

        got = darr["T":"A", [1, 2]]
        assert_allclose(got.array, numpy.array([[0.1, 0.2], [0.7, 0.1]]))

        # make sure we cope with keys that are int's
        nums = list(range(1, 5))
        darr = DictArrayTemplate(nums, nums).wrap(
            [
                [0.7, 0.1, 0.2, 0.3],
                [0.1, 0.7, 0.1, 0.3],
                [0.3, 0.2, 0.6, 0.3],
                [0.4, 0.1, 0.1, 0.7],
            ]
        )
        got = darr[[1, 2], [1, 2]]
        assert_allclose(got.array, numpy.array([[0.7, 0.1], [0.2, 0.6]]))

    def test_add(self):
        """can add compatible dict arrays"""
        data = numpy.array([[7, 1], [1, 7]])
        darr1 = DictArrayTemplate(list("AB"), list("CD")).wrap(data)
        darr2 = DictArrayTemplate(list("AB"), list("CD")).wrap(data)
        darr3 = darr1 + darr2
        assert_allclose(darr3.array, 2 * data)
        self.assertEqual(darr3.template.names, darr1.template.names)
        # must be correct type
        with self.assertRaises(TypeError):
            darr1 + data

        # must be equal dimensions
        with self.assertRaises(ValueError):
            darr1 + DictArrayTemplate(list("CD"), list("AB")).wrap(data)