# coding: utf-8

# ------------------------------------------------------------------------------
# Copyright (C) 2018-2019 Maximilian Stahlberg
#
# This file is part of PICOS.
#
# PICOS is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# PICOS is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ------------------------------------------------------------------------------

"""Affine constraint types."""

from collections import namedtuple

from .. import glyphs
from ..apidoc import api_end, api_start
from ..caching import cached_property
from .constraint import Constraint, ConstraintConversion

_API_START = api_start(globals())
# -------------------------------


class AffineConstraint(Constraint):
    """An equality or inequality between two affine expressions."""

    def __init__(self, lhs, relation, rhs, customString=None):
        """Construct an :class:`AffineConstraint`.

        :param ~picos.expressions.AffineExpression lhs:
            Left hand side expression.
        :param str relation:
            Constraint relation symbol.
        :param ~picos.expressions.AffineExpression rhs:
            Right hand side expression.
        :param str customString:
            Optional string description.
        """
        from ..expressions import AffineExpression

        assert isinstance(lhs, AffineExpression)
        assert isinstance(rhs, AffineExpression)
        assert relation in self.LE + self.GE + self.EQ
        assert lhs.size == rhs.size

        self.lhs      = lhs
        self.rhs      = rhs
        self.relation = relation

        super(AffineConstraint, self).__init__(
            "Affine", customString, printSize=True)

    Subtype = namedtuple("Subtype", ("dim", "eq"))

    def _subtype(self):
        return self.Subtype(len(self.lhs), self.relation == self.EQ)

    @property
    def smaller(self):
        """Smaller-or-equal side of the constraint.

        The smaller-or-equal side expression in case of an inequality, otherwise
        the left hand side.
        """
        return self.rhs if self.relation == self.GE else self.lhs

    @property
    def greater(self):
        """Greater-or-equal side of the constraint.

        The greater-or-equal side expression in case of an inequality, otherwise
        the right hand side.
        """
        return self.lhs if self.relation == self.GE else self.rhs

    @cached_property
    def le0(self):
        """Expression constrained to be lower than or equal to zero.

        The expression posed to be less than or equal to zero in case of an
        inequality, otherwise the left hand side minus the right hand side.
        """
        if self.relation == self.GE:
            return self.rhs - self.lhs
        else:
            return self.lhs - self.rhs

    @cached_property
    def ge0(self):
        """Expression constrained to be greater than or equal to zero.

        The expression posed to be greater than or equal to zero in case of an
        inequality, otherwise the right hand side minus the left hand side.
        """
        if self.relation == self.GE:
            return self.lhs - self.rhs
        else:
            return self.rhs - self.lhs

    @classmethod
    def _cost(cls, subtype):
        return subtype.dim

    def _expression_names(self):
        yield "lhs"
        yield "rhs"

    def _str(self):
        if self.relation == self.LE:
            return glyphs.le(self.lhs.string, self.rhs.string)
        elif self.relation == self.GE:
            return glyphs.ge(self.lhs.string, self.rhs.string)
        else:
            return glyphs.eq(self.lhs.string, self.rhs.string)

    def _get_size(self):
        return self.lhs.size

    def _get_slack(self):
        if self.relation == self.LE:
            delta = self.rhs.value - self.lhs.value
        else:
            delta = self.lhs.value - self.rhs.value

        return -abs(delta) if self.relation == self.EQ else delta

    def bounded_linear_form(self):
        """Bounded linear form of the constraint.

        Separates the constraint into a linear function on the left hand side
        and a constant bound on the right hand side.

        :returns: A pair ``(linear, bound)`` where ``linear`` is a pure linear
            expression and ``bound`` is a constant expression.
        """
        linear = self.lhs - self.rhs
        bound  = -linear.cst
        linear = linear + bound

        return (linear, bound)

    def sparse_Ab_rows(self, varOffsetMap, indexFunction=None):
        """Sparse representation of the constraint's bounded linear form.

        A sparse list representation of the constraint, given a mapping of PICOS
        variables to column offsets (or alternatively given an index function).

        The constraint is brought into a bounded linear form A • b, where • is
        one of ≤, ≥, or =, depending on the constraint relation, and the rows
        returned correspond to the matrix [A|b].

        :param dict varOffsetMap: Maps variables or variable start indices to
            column offsets.
        :param indexFunction: Instead of adding the local variable index to the
            value returned by varOffsetMap, use the return value of this
            function, that takes as argument the variable and its local index,
            as the "column index", which need not be an integer. When this
            parameter is passed, the parameter varOffsetMap is ignored.
        :returns: A list of triples (J, V, c) where J contains column indices
            (representing scalar variables), V contains coefficients for each
            column index and c is a constant term. Each entry of the list
            represents a row in a constraint matrix.
        """
        lhs  = self.lhs - self.rhs
        rows = lhs.sparse_rows(varOffsetMap, indexFunction=indexFunction)

        for localConIndex in range(len(lhs)):
            rows[localConIndex][2] = -rows[localConIndex][2]

        return rows


class ComplexAffineConstraint(Constraint):
    """An equality between affine expressions, at least one being complex."""

    class RealConversion(ConstraintConversion):
        """Complex affine equality to real affine equality conversion."""

        @classmethod
        def predict(cls, subtype, options):
            """Implement :meth:`~.constraint.ConstraintConversion.predict`."""
            yield ("con",
                AffineConstraint.make_type(dim=2*subtype.dim, eq=True), 1)

        @classmethod
        def convert(cls, con, options):
            """Implement :meth:`~.constraint.ConstraintConversion.convert`."""
            from ..modeling import Problem

            P = Problem()
            P.add_constraint((con.lhs.real // con.lhs.imag)
                == (con.rhs.real // con.rhs.imag))

            return P

        @classmethod
        def dual(cls, auxVarPrimals, auxConDuals, options):
            """Implement :meth:`~.constraint.ConstraintConversion.dual`."""
            assert len(auxConDuals) == 1

            auxConDual = auxConDuals[0]
            if auxConDual is None:
                return None
            else:
                n = auxConDual.size[0] // 2
                return auxConDual[:n, :] + 1j*auxConDual[n:, :]

    def __init__(self, lhs, rhs, customString=None):
        """Construct a :class:`ComplexAffineConstraint`.

        :param ~picos.expressions.AffineExpression lhs:
            Left hand side expression.
        :param ~picos.expressions.AffineExpression rhs:
            Right hand side expression.
        :param str customString:
            Optional string description.
        """
        from ..expressions import ComplexAffineExpression

        assert isinstance(lhs, ComplexAffineExpression)
        assert isinstance(rhs, ComplexAffineExpression)
        assert lhs.size == rhs.size

        self.lhs = lhs
        self.rhs = rhs

        super(ComplexAffineConstraint, self).__init__(
            "Complex Equality", customString, printSize=True)

    Subtype = namedtuple("Subtype", ("dim",))

    def _subtype(self):
        return self.Subtype(len(self.lhs))

    @classmethod
    def _cost(cls, subtype):
        return 2*subtype.dim

    def _expression_names(self):
        yield "lhs"
        yield "rhs"

    def _str(self):
        return glyphs.eq(self.lhs.string, self.rhs.string)

    def _get_size(self):
        return self.lhs.size

    def _get_slack(self):
        return -abs(self.lhs.value - self.rhs.value)


# --------------------------------------
__all__ = api_end(_API_START, globals())
