path: root/polymatrix/__init__.py

import collections
import dataclasses
import itertools
import math
import typing
import numpy as np
import scipy.sparse
import sympy

from polymatrix.expression.expression import Expression
from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
from polymatrix.expression.mixins.expressionmixin import ExpressionMixin
from polymatrix.expression.mixins.parametrizeexprmixin import ParametrizeExprMixin
from polymatrix.expressionstate.expressionstate import ExpressionState
from polymatrix.expression.init.initblockdiagexpr import init_block_diag_expr
from polymatrix.expression.init.initexpression import init_expression
from polymatrix.expression.init.initeyeexpr import init_eye_expr
from polymatrix.expression.init.initfromsympyexpr import init_from_sympy_expr
from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
from polymatrix.expression.init.initvstackexpr import init_v_stack_expr
from polymatrix.polymatrix.polymatrix import PolyMatrix
from polymatrix.expression.utils.getvariableindices import get_variable_indices, get_variable_indices_from_variable
from polymatrix.statemonad.init.initstatemonad import init_state_monad
from polymatrix.statemonad.mixins.statemonadmixin import StateMonadMixin
from polymatrix.expression.utils.monomialtoindex import monomial_to_index
from polymatrix.expressionstate.init.initexpressionstate import init_expression_state as original_init_expression_state
from polymatrix.statemonad.statemonad import StateMonad


def init_expression_state():
    return original_init_expression_state()

def from_sympy(
    data: tuple[tuple[float]],
):
    return init_expression(
        init_from_sympy_expr(data)
    )

def from_state_monad(
    data: StateMonad,
):
    return init_expression(
        data.flat_map(lambda inner_data: init_from_sympy_expr(inner_data)),
    )

def from_polymatrix(
    polymatrix: PolyMatrix,
):
    return init_expression(
        init_from_terms_expr(polymatrix)
    )

def from_(
    data: tuple[tuple[float]],
):
    return from_sympy(data)

def v_stack(
    expressions: tuple[Expression],
):
    return init_expression(
        init_v_stack_expr(expressions)
    )

def h_stack(
    expressions: tuple[Expression],
):
    return init_expression(
        init_v_stack_expr(tuple(expr.T for expr in expressions))
    ).T

def block_diag(
    expressions: tuple[Expression],
):
    return init_expression(
        init_block_diag_expr(expressions)
    )

def eye(
    variable: tuple[Expression],
):
    return init_expression(
        init_eye_expr(variable=variable)
    )

def kkt_equality(
    variables: Expression,
    equality: Expression = None,
):

    self_variables = variables
    self_equality = equality

    def func(state: ExpressionState):

        state, equality = self_equality.apply(state=state)

        state, equality_der = self_equality.diff(
            self_variables, 
            introduce_derivatives=True,
        ).apply(state)

        def acc_nu_variables(acc, v):
            state, nu_variables = acc

            nu_variable = state.n_param
            state = state.register(n_param=1)
            
            return state, nu_variables + [nu_variable]

        *_, (state, nu_variables) = tuple(itertools.accumulate(
            range(equality.shape[0]), 
            acc_nu_variables, 
            initial=(state, []),
        ))

        terms = {}

        for row in range(equality_der.shape[1]):

            monomial_terms = collections.defaultdict(float)

            for eq_idx, nu_variable in enumerate(nu_variables):

                underlying_terms = equality_der.get_poly(eq_idx, row)
                if underlying_terms is None:
                    continue

                for monomial, value in underlying_terms.items():
                    new_monomial = monomial + (nu_variable,)

                    monomial_terms[new_monomial] += value

            # terms[row, 0] = dict(monomial_terms)
            terms[row, 0] = monomial_terms

        cost_expr = init_expression(init_from_terms_expr(
            terms=terms,
            shape=(equality_der.shape[1], 1),
        ))

        nu_terms = {}

        for eq_idx, nu_variable in enumerate(nu_variables):
            nu_terms[eq_idx, 0] = {(nu_variable,): 1}

        nu_expr = init_expression(init_from_terms_expr(
            terms=nu_terms,
            shape=(len(nu_variables), 1),
        ))

        return state, (nu_expr, cost_expr)   

    return init_state_monad(func)

def kkt_inequality(
    variables: Expression,
    inequality: Expression = None,
):

    self_variables = variables
    self_inequality = inequality

    def func(state: ExpressionState):

        state, inequality = self_inequality.apply(state=state)

        state, inequality_der = self_inequality.diff(
            self_variables, 
            introduce_derivatives=True,
        ).apply(state)

        def acc_lambda_variables(acc, v):
            state, lambda_variables = acc

            lambda_variable = state.n_param
            state = state.register(n_param=3)
            
            return state, lambda_variables + [lambda_variable]

        *_, (state, lambda_variables) = tuple(itertools.accumulate(
            range(inequality.shape[0]), 
            acc_lambda_variables, 
            initial=(state, []),
        ))

        terms = {}

        for row in range(inequality_der.shape[1]):

            monomial_terms = collections.defaultdict(float)

            for inequality_idx, lambda_variable in enumerate(lambda_variables):

                polynomial = inequality_der.get_poly(inequality_idx, row)
                if polynomial is None:
                    continue

                for monomial, value in polynomial.items():
                    new_monomial = monomial + (lambda_variable,)

                    monomial_terms[new_monomial] += value

            # terms[row, 0] = dict(monomial_terms)
            terms[row, 0] = monomial_terms

        cost_expr = init_expression(init_from_terms_expr(
            terms=terms,
            shape=(inequality_der.shape[1], 1),
        ))

        # inequality, dual feasibility, complementary slackness
        # -----------------------------------------------------

        inequality_terms = {}
        feasibility_terms = {}
        complementary_terms = {}

        for inequality_idx, lambda_variable in enumerate(lambda_variables):
            r_lambda = lambda_variable + 1
            r_inequality = lambda_variable + 2

            polynomial = inequality.get_poly(inequality_idx, 0)
            if polynomial is None:
                continue

            # f(x) <= -0.01
            inequality_terms[inequality_idx, 0] = polynomial | {(r_inequality, r_inequality): 1}

            # dual feasibility, lambda >= 0
            feasibility_terms[inequality_idx, 0] = {(lambda_variable,): 1, (r_lambda, r_lambda): -1}

            # complementary slackness
            complementary_terms[inequality_idx, 0] = {(r_lambda, r_inequality): 1}

        inequality_expr = init_expression(init_from_terms_expr(
            terms=inequality_terms,
            shape=(len(lambda_variables), 1),
        ))

        feasibility_expr = init_expression(init_from_terms_expr(
            terms=feasibility_terms,
            shape=(len(lambda_variables), 1),
        ))

        complementary_expr = init_expression(init_from_terms_expr(
            terms=complementary_terms,
            shape=(len(lambda_variables), 1),
        ))

        # lambda expression
        # -----------------

        terms = {}
        for inequality_idx, lambda_variable in enumerate(lambda_variables):
            terms[inequality_idx, 0] = {(lambda_variable,): 1}

        lambda_expr = init_expression(init_from_terms_expr(
            terms=terms,
            shape=(len(lambda_variables), 1),
        ))

        return state, (lambda_expr, cost_expr, inequality_expr, feasibility_expr, complementary_expr)   

    return init_state_monad(func)

def rows(
    expr: Expression,
) -> StateMonadMixin[ExpressionState, np.ndarray]:

    def func(state: ExpressionState):
        state, underlying = expr.apply(state)

        def gen_row_terms():
            for row in range(underlying.shape[0]):

                terms = {}

                for col in range(underlying.shape[1]):
                    polynomial = underlying.get_poly(row, col)
                    if polynomial == None:
                        continue

                    terms[0, col] = polynomial

                yield init_expression(underlying=init_from_terms_expr(
                    terms=terms,
                    shape=(1, underlying.shape[1])
                ))

        row_terms = tuple(gen_row_terms())

        return state, row_terms

    return init_state_monad(func)

def shape(
    expr: Expression,
) -> StateMonadMixin[ExpressionState, tuple[int, ...]]:
    def func(state: ExpressionState):
        state, polymatrix = expr.apply(state)

        return state, polymatrix.shape

    return init_state_monad(func)


@dataclasses.dataclass
class MatrixBuffer:
    data: dict[int, ...]
    n_row: int
    n_param: int

    def get_max_degree(self):
        return max(degree for degree in self.data.keys())

    def add_buffer(self, index: int):
        if index <= 1:
            buffer = np.zeros((self.n_row, self.n_param**index), dtype=np.double)

        else:
            buffer = scipy.sparse.dok_array((self.n_row, self.n_param**index), dtype=np.double)
            
        self.data[index] = buffer

    def add(self, row: int, col: int, index: int, value: float):
        if index not in self.data:
            self.add_buffer(index)

        self.data[index][row, col] = value

    def __getitem__(self, key):
        if key not in self.data:
            self.add_buffer(key)

        return self.data[key]


@dataclasses.dataclass
class MatrixRepresentations:
    data: tuple[MatrixBuffer, ...]
    aux_data: typing.Optional[MatrixBuffer]
    variable_mapping: tuple[int, ...]
    state: ExpressionState

    def merge_matrix_equations(self):
        def gen_matrices(index: int):
            for equations in self.data:
                if index < len(equations):
                    yield equations[index]

            if index < len(self.aux_data):
                yield self.aux_data[index]

        indices = set(key for equations in self.data + (self.aux_data,) for key in equations.keys())

        def gen_matrices():
            for index in indices:
                if index <= 1:
                    yield index, np.vstack(tuple(gen_matrices(index)))
                else:
                    yield index, scipy.sparse.vstack(tuple(gen_matrices(index)))

        return dict(gen_matrices())

    def get_value(self, variable, value):
        variable_indices = get_variable_indices_from_variable(self.state, variable)[1]

        def gen_value_index():
            for variable_index in variable_indices:
                try:
                    yield self.variable_mapping.index(variable_index)
                except ValueError:
                    raise ValueError(f'{variable_index} not found in {self.variable_mapping}')

        value_index = list(gen_value_index())

        return value[value_index]

    def set_value(self, variable, value):
        variable_indices = get_variable_indices_from_variable(self.state, variable)[1]
        value_index = list(self.variable_mapping.index(variable_index) for variable_index in variable_indices)
        vec = np.zeros(len(self.variable_mapping))
        vec[value_index] = value
        return vec

    def get_func(self, eq_idx: int):
        equations = self.data[eq_idx].data
        max_idx = max(equations.keys())

        if 2 <= max_idx:
            def func(x: np.ndarray) -> np.ndarray:
                if isinstance(x, tuple) or isinstance(x, list):
                    x = np.array(x).reshape(-1, 1)

                def acc_x_powers(acc, _):
                    next = (acc @ x.T).reshape(-1, 1)
                    return next

                x_powers = tuple(itertools.accumulate(
                    range(max_idx - 1),
                    acc_x_powers,
                    initial=x,
                ))[1:]

                def gen_value():
                    for idx, equation in equations.items():
                        if idx == 0:
                            yield equation

                        elif idx == 1:
                            yield equation @ x

                        else:
                            yield equation @ x_powers[idx-2]

                return sum(gen_value())

        else:
            def func(x: np.ndarray) -> np.ndarray:
                def gen_value():
                    for idx, equation in equations.items():
                        if idx == 0:
                            yield equation

                        else:
                            yield equation @ x

                return sum(gen_value())

        return func

def to_matrix_repr(
    expressions: Expression | tuple[Expression],
    variables: Expression,
) -> StateMonadMixin[ExpressionState, tuple[tuple[tuple[np.ndarray, ...], ...], tuple[int, ...]]]:

    if isinstance(expressions, Expression):
        expressions = (expressions,)

    assert isinstance(variables, ExpressionBaseMixin), f'{variables=}'

    def func(state: ExpressionState):

        def acc_underlying_application(acc, v):
            state, underlying_list = acc

            state, underlying = v.apply(state)

            assert underlying.shape[1] == 1, f'{underlying.shape[1]=} is not 1'
            
            return state, underlying_list + (underlying,)

        *_, (state, underlying_list) = tuple(itertools.accumulate(
            expressions, 
            acc_underlying_application, 
            initial=(state, tuple()),
        ))

        state, ordered_variable_index = get_variable_indices_from_variable(state, variables)

        assert len(ordered_variable_index) == len(set(ordered_variable_index)), f'{ordered_variable_index=} contains repeated variables'

        variable_index_map = {old: new for new, old in enumerate(ordered_variable_index)}

        n_param = len(ordered_variable_index)

        def gen_underlying_matrices():
            for underlying in underlying_list:

                n_row = underlying.shape[0]

                buffer = MatrixBuffer(
                    data={},
                    n_row=n_row,
                    n_param=n_param,
                )

                for row in range(n_row):
                    underlying_terms = underlying.get_poly(row, 0)
                    if underlying_terms is None:
                        continue
                    
                    for monomial, value in underlying_terms.items():
                        
                        def gen_new_monomial():
                            for var, count in monomial:
                                try:
                                    index = variable_index_map[var]
                                except KeyError:
                                    raise KeyError(f'{var=} ({state.get_key_from_offset(var)}) is incompatible with {variable_index_map=}')

                                for _ in range(count):
                                    yield index

                        new_monomial = tuple(gen_new_monomial())

                        cols = monomial_to_index(n_param, new_monomial)

                        col_value = value / len(cols)

                        for col in cols:
                            buffer.add(row, col, sum(count for _, count in monomial), col_value)
                
                yield buffer

        underlying_matrices = tuple(gen_underlying_matrices())

        def gen_auxillary_equations():
            for key, monomial_terms in state.auxillary_equations.items():
                if key in ordered_variable_index:
                    yield key, monomial_terms

        auxillary_equations = tuple(gen_auxillary_equations())

        n_row = len(auxillary_equations)

        if n_row == 0:
            auxillary_matrix_equations = None

        else:
            buffer = MatrixBuffer(
                data={},
                n_row=n_row,
                n_param=n_param,
            )

            for row, (key, monomial_terms) in enumerate(auxillary_equations):
                for monomial, value in monomial_terms.items():
                    new_monomial = tuple(variable_index_map[var] for var, count in monomial for _ in range(count))

                    cols = monomial_to_index(n_param, new_monomial)

                    col_value = value / len(cols)

                    for col in cols:
                        buffer.add(row, col, sum(count for _, count in monomial), col_value)

            auxillary_matrix_equations = buffer

        result = MatrixRepresentations(
            data=underlying_matrices,
            aux_data=auxillary_matrix_equations,
            variable_mapping=ordered_variable_index,
            state=state,
        )

        return state, result

    return init_state_monad(func)

def to_constant_repr(
    expr: Expression,
    assert_constant: bool = True,
) -> StateMonadMixin[ExpressionState, np.ndarray]:

    def func(state: ExpressionState):
        state, underlying = expr.apply(state)

        A = np.zeros(underlying.shape, dtype=np.double)

        for (row, col), polynomial in underlying.gen_terms():
            for monomial, value in polynomial.items():
                if len(monomial) == 0:
                    A[row, col] = value
                
                elif assert_constant:
                    raise Exception(f'non-constant term {monomial=}')

        return state, A

    return init_state_monad(func)


def degrees(
    expr: Expression,
    variables: Expression,
) -> StateMonadMixin[ExpressionState, np.ndarray]:

    def func(state: ExpressionState):
        state, underlying = expr.apply(state)
        state, variable_indices = get_variable_indices_from_variable(state, variables)

        def gen_rows():
            for row in range(underlying.shape[0]):
                def gen_cols():
                    for col in range(underlying.shape[1]):
                        
                        def gen_degrees():
                            polynomial = underlying.get_poly(row, col)

                            if polynomial is None:
                                yield 0

                            else:
                                for monomial, _ in polynomial.items():
                                    yield sum(count for var, count in monomial if var in variable_indices)

                        yield tuple(set(gen_degrees()))

                yield tuple(gen_cols())

        return state, tuple(gen_rows())

    return init_state_monad(func)


def to_sympy_repr(
    expr: Expression,
) -> StateMonadMixin[ExpressionState, sympy.Expr]:

    def func(state: ExpressionState):
        state, underlying = expr.apply(state)

        A = np.zeros(underlying.shape, dtype=object)

        for (row, col), polynomial in underlying.gen_terms():

            sympy_polynomial = 0

            for monomial, value in polynomial.items():
                sympy_monomial = 1

                for offset, count in monomial:

                    variable = state.get_key_from_offset(offset)
                    # def get_variable_from_offset(offset: int):
                    #     for variable, (start, end) in state.offset_dict.items():
                    #         if start <= offset < end:
                                # assert end - start == 1, f'{start=}, {end=}, {variable=}'

                    if isinstance(variable, sympy.core.symbol.Symbol):
                        variable_name = variable.name
                    elif isinstance(variable, ParametrizeExprMixin):
                        variable_name = variable.name
                    elif isinstance(variable, str):
                        variable_name = variable
                    else:
                        raise Exception(f'{variable=}')

                    start, end = state.offset_dict[variable]

                    if end - start == 1:
                        sympy_var = sympy.Symbol(variable_name)
                    else:
                        sympy_var = sympy.Symbol(f'{variable_name}_{offset - start + 1}')

                    # var = get_variable_from_offset(offset)
                    sympy_monomial *= sympy_var**count

                sympy_polynomial += value * sympy_monomial

            A[row, col] = sympy_polynomial

        return state, A

    return init_state_monad(func)