import unittest

from polymatrix.polystruct import init_equation, init_poly_matrix, init_poly_vector
from polymatrix.utils import variable_to_index

class TestPolyMatrix(unittest.TestCase):
    @staticmethod
    def assert_term_in_eq(terms, degree, eq_idx, row_idx, value, monoms=None):
        if monoms is None:
            monoms = tuple()

        assert degree in terms, f'could not find {degree} in {terms}'
        degree_terms = terms[degree]

        key = eq_idx, monoms
        assert key in degree_terms, f'could not find {key} in {degree_terms}'
        eq_terms = degree_terms[key]

        assert row_idx in eq_terms, f'could not find {row_idx} in {eq_terms}'
        assert eq_terms[row_idx] == value, f'value {eq_terms[row_idx]} and {value} do not match'

    def test_param_matrix_param_vector_degree_0(self):
        """
        param = [a11 a21 a31 a41 v11 v21]
        """

        mat = init_poly_matrix(degrees=(0,))
        vec = init_poly_vector(degrees=(0,))
        n_param = 6

        eq = init_equation(
            terms = [(mat, vec)],
            n_var = 2,
        )

        terms, offset_dict = list(eq.create())

        # a11 v11
        self.assert_term_in_eq(
            terms = terms,
            degree = 2,
            eq_idx = 0,
            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
            value = 1,
        )

        # a21 v11
        self.assert_term_in_eq(
            terms = terms,
            degree = 2,
            eq_idx = 1,
            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)])),
            value = 1,
        )

    def test_param_matrix_param_vector_degree_01(self):
        """
        param = [a11 a21 a31 a41 v011 v021 v111 v112 v121 v122]
        """

        mat = init_poly_matrix(degrees=(0,))
        vec = init_poly_vector(degrees=(0,1))
        n_param = 10

        eq = init_equation(
            terms = [(mat, vec)],
            n_var = 2,
        )

        terms, offset_dict = list(eq.create())

        # a11 v011
        self.assert_term_in_eq(
            terms = terms,
            degree = 2,
            eq_idx = 0,
            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
            value = 1,
        )

        # a11 v011
        self.assert_term_in_eq(
            terms = terms,
            degree = 2,
            eq_idx = 0,
            monoms=(0,),
            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 1)])),
            value = 1,
        )

    def test_param_matrix_const_vector_degree_0(self):
        """
        param = [a11 a21 a31 a41]
        """

        mat = init_poly_matrix(degrees=(0,))
        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})

        eq = init_equation(
            terms = [(mat, vec)],
            n_var = 2,
        )

        terms, offset_dict = list(eq.create())

        # a11
        self.assert_term_in_eq(
            terms = terms,
            degree = 1,
            eq_idx = 0,
            row_idx = offset_dict[(mat, 0)],
            value = 1,
        )

    def test_const_matrix_const_vector_degree_0(self):
        """
        param = [a11 a21 a31 a41]
        """

        mat = init_poly_matrix(subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}})
        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})

        eq = init_equation(
            terms = [(mat, vec)],
            n_var = 2,
        )

        terms, offset_dict = list(eq.create())

        self.assert_term_in_eq(
            terms = terms,
            degree = 0,
            eq_idx = 0,
            row_idx = 0,
            value = 2,
        )
    
    def test_param_matrix_const_vector_skew_symmetric(self):
        """
        param = [a11 a21 a31 a41]
        """

        def skew_symmetric(idx1, idx2):
            if idx1 == idx2:
                return idx1, idx2, 0
            elif idx2 < idx1:
                return idx2, idx1, -1

        mat = init_poly_matrix(
            degrees=(0,), 
            re_index_func=skew_symmetric,
        )
        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})

        eq = init_equation(
            terms = [(mat, vec)],
            n_var = 2,
        )

        terms, offset_dict = list(eq.create())

        self.assert_term_in_eq(
            terms = terms,
            degree = 1,
            eq_idx = 0,
            row_idx = offset_dict[(mat, 0)] + 1,
            value = 1,
        )

        self.assert_term_in_eq(
            terms = terms,
            degree = 1,
            eq_idx = 1,
            row_idx = offset_dict[(mat, 0)] + 1,
            value = -1,
        )