import unittest
from polymatrix.init.initoptimization import init_optimization
from polymatrix.init.initpolymatrix import init_poly_matrix
from polymatrix.optimization import Optimization
from polymatrix.polymatrix import PolyMatrix

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

    #     assert degree in result.data, f'could not find {degree} in {result.data}'
    #     degree_terms = result.data[degree]

    #     key = eq_idx, monoms, row_idx
    #     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 degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'

    # @staticmethod
    # def assert_term_in_eq(result, degree, row_idx, monoms, value):

    #     assert degree in result.data, f'could not find {degree} in {result.data}'
    #     degree_terms = result.data[degree]

    #     key = row_idx, variable_to_index(result.n_param, monoms)
    #     assert key in degree_terms, f'could not find {key} in {degree_terms}'

    #     assert degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'

    @staticmethod
    def assert_term_in_eq(
        problem: Optimization,
        poly_row: int, 
        p_monomial: tuple[tuple[PolyMatrix, int, int], ...], 
        value: float,
        x_monomial: tuple[int, ...] = None,
    ):
        if x_monomial is None:
            x_monomial = tuple()

        offset_dict = problem.state.offset_dict
        p_monomial = tuple(offset_dict[(polymat, degree)][0] + offset for polymat, degree, offset in p_monomial)

        equality_constraint = problem.equality_constraints[0]

        key = (poly_row, x_monomial, p_monomial)
        assert key in equality_constraint, f'could not find {key} in {equality_constraint}'

        assert equality_constraint[key] == value, f'value {equality_constraint[key]} and {value} do not match'

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

        n_var = 2

        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
        vec = init_poly_matrix(name='vec', degrees=(0,), shape=(n_var, 1))

        problem = init_optimization(
            n_var=n_var,
        ).add_equality_constraints(
            expr=[(mat, vec)],
        )

        # a11 v11
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            p_monomial = ((mat, 0, 0), (vec, 0, 0)),
            value = 1,
        )

        # a12 v21
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            p_monomial = ((mat, 0, 2), (vec, 0, 1)),
            value = 1,
        )

        # a21 v11
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 1,
            p_monomial = ((mat, 0, 1), (vec, 0, 0)),
            value = 1,
        )

        # a22 v21
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 1,
            p_monomial = ((mat, 0, 3), (vec, 0, 1)),
            value = 1,
        )

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

        n_var = 2

        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
        vec = init_poly_matrix(name='vec', degrees=(0, 1), shape=(n_var, 1))

        problem = init_optimization(
            n_var=n_var,
        ).add_equality_constraints(
            expr=[(mat, vec)],
        )

        # a11 v011
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            p_monomial = ((mat, 0, 0), (vec, 0, 0)),
            value = 1,
        )

        # a11 v011
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            x_monomial=(0,),
            p_monomial = ((mat, 0, 0), (vec, 1, 0)),
            value = 1,
        )

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

        n_var = 2

        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))

        problem = init_optimization(
            n_var=n_var,
        ).add_equality_constraints(
            expr=[(mat, vec)],
        )


        # a11
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            p_monomial = ((mat, 0, 0),),
            value = 1,
        )

        # a21
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 1,
            p_monomial = ((mat, 0, 1),),
            value = 1,
        )

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

        n_var = 2
        
        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
        vec = init_poly_matrix(name='vec', subs={1: {(0, 0, 0): 1, (0, 0, 1): 0, (1, 0, 0): 0, (1, 0, 1): 1}}, shape=(n_var, 1))

        problem = init_optimization(
            n_var=n_var,
        ).add_equality_constraints(
            expr=[(mat, vec)],
        )

        # a11
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            x_monomial=(0,),
            p_monomial = ((mat, 0, 0),),
            value = 1,
        )

        # a12
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            x_monomial=(1,),
            p_monomial = ((mat, 0, 2),),
            value = 1,
        )

        # a21
        self.assert_term_in_eq(
            problem = problem,
            poly_row = 1,
            x_monomial=(0,),
            p_monomial = ((mat, 0, 1),),
            value = 1,
        )

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

        n_var = 2

        mat = init_poly_matrix(name='mat', subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, shape=(n_var, n_var))
        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))

        problem = init_optimization(
            n_var=n_var,
        ).add_equality_constraints(
            expr=[(mat, vec)],
        )

        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            p_monomial = tuple(),
            value = 2,
        )
    
    def test_skew_symmetric_param_matrix_const_vector(self):
        """
        param = [a11 a21 a31 a41]
        """

        def skew_symmetric(degree, poly_row, poly_col, monom):
            if poly_row == poly_col:
                return poly_row, poly_col, monom, 0
            elif poly_col < poly_row:
                return poly_col, poly_row, monom, -1

        n_var = 2

        mat = init_poly_matrix(
            name='mat', 
            degrees=(0,), 
            re_index=skew_symmetric,
            shape=(n_var, n_var),
        )
        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))

        problem = init_optimization(
            n_var=n_var,
        ).add_equality_constraints(
            expr=[(mat, vec)],
        )

        self.assert_term_in_eq(
            problem = problem,
            poly_row = 0,
            p_monomial = ((mat, 0, 2),),
            value = 1,
        )

        self.assert_term_in_eq(
            problem = problem,
            poly_row = 1,
            p_monomial = ((mat, 0, 2),),
            value = -1,
        )

    def test_const_matrix_d0_param_gradient_vector_d1(self):
        """
        param = [v11 v12 v21 v22]
        """

        def gradient(degree, p_row, p_col, monom):
            if degree == 1:
                factor = sum(p_row==e for e in monom) + 1

                if monom[-1] < p_row:
                    n_p_row = monom[-1]
                    n_monom = sorted(monom + (p_row,), reverse=True)

                if p_row <= monom[-1]:
                    n_p_row = p_row
                    n_monom = monom

                return n_p_row, p_col, n_monom, factor

        n_var = 2

        mat = init_poly_matrix(
            name='mat', 
            subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, 
            shape=(n_var, n_var),
        )
        vec = init_poly_matrix(
            name='vec', 
            degrees=(1,),
            re_index=gradient,
            shape=(n_var, 1),
        )

        problem = init_optimization(
            n_var=n_var,
        ).add_equality_constraints(
            expr=[(mat, vec)],
        )

        self.assert_term_in_eq(
            problem = problem,
            x_monomial=(0,),
            poly_row = 0,
            p_monomial = ((vec, 1, 0),),
            value = 2,
        )

        self.assert_term_in_eq(
            problem = problem,
            x_monomial=(1,),
            poly_row = 0,
            p_monomial = ((vec, 1, 2),),
            value = 1,
        )

        self.assert_term_in_eq(
            problem = problem,
            x_monomial=(0,),
            poly_row = 1,
            p_monomial = ((vec, 1, 2),),
            value = 1,
        )