# Copyright 2017 Francesco Ceccon
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""FBBT bounds tightening rules."""
from suspect.expression import UnaryFunctionType
from suspect.interfaces import Rule, UnaryFunctionRule
from suspect.interval import Interval
from suspect.math import almosteq  # pylint: disable=no-name-in-module
from suspect.pyomo.expressions import nonpyomo_leaf_types

MAX_EXPR_CHILDREN = 1000


class ConstraintRule(Rule):
    """Return new bounds for constraint."""
    def apply(self, expr, _bounds):
        bounds = Interval(expr.lower_bound, expr.upper_bound)
        return [bounds]


class SumRule(Rule):
    """Return new bounds for sum."""
    def __init__(self, max_expr_children=None):
        if max_expr_children is None:
            max_expr_children = MAX_EXPR_CHILDREN
        self.max_expr_children = max_expr_children

    def apply(self, expr, bounds):
        if len(expr.args) > self.max_expr_children: # pragma: no cover
            return None
        expr_bound = bounds[expr]
        children_bounds = [
            self._child_bounds(child_idx, child, expr, expr_bound, bounds)
            for child_idx, child in enumerate(expr.args)
        ]
        return children_bounds

    def _child_bounds(self, child_idx, _child, expr, expr_bound, bounds):
        siblings_bound = sum(
            bounds[c]
            for i, c in enumerate(expr.args)
            if i != child_idx
        )
        return expr_bound - siblings_bound


class LinearRule(Rule):
    """Return new bounds for linear expressions."""
    def __init__(self, max_expr_children=None):
        if max_expr_children is None:
            max_expr_children = MAX_EXPR_CHILDREN
        self.max_expr_children = max_expr_children

    def apply(self, expr, bounds):
        if expr.nargs() > self.max_expr_children:  # pragma: no cover
            return None
        expr_bound = bounds[expr]
        const = expr.constant
        children_bounds = [
            self._child_bounds(child_idx, coef, child, expr, const, expr_bound, bounds)
            for child_idx, (coef, child) in enumerate(zip(expr.linear_coefs, expr.linear_vars))
        ]
        return children_bounds

    def _child_bounds(self, child_idx, coef, _child, expr, const, expr_bound, bounds):
        siblings_bound = sum(
            bounds[s] * c
            for i, (c, s) in enumerate(zip(expr.linear_coefs, expr.linear_vars))
            if i != child_idx
        ) + const
        return (expr_bound - siblings_bound) / coef


class QuadraticRule(Rule):
    """Return new bounds for quadratic expressions."""
    def __init__(self, max_expr_children=None):
        if max_expr_children is None:
            max_expr_children = MAX_EXPR_CHILDREN
        self.max_expr_children = max_expr_children

    def apply(self, expr, bounds):
        expr_bound = bounds[expr]
        child_bounds = {}
        if len(expr.terms) > self.max_expr_children:
            return None

        terms = expr.terms
        for term_idx, term in enumerate(terms):
            var1 = term.var1
            var2 = term.var2
            siblings_bound = sum(
                self._term_bound(t, bounds)
                for i, t in enumerate(terms) if i != term_idx
            )
            term_bound = (expr_bound - siblings_bound) / term.coefficient
            if var1 is var2:
                term_bound = term_bound.intersect(Interval(0, None))
                upper_bound = term_bound.sqrt().upper_bound
                new_bound = Interval(-upper_bound, upper_bound)

                if id(var1) in child_bounds:
                    existing = child_bounds[id(var1)]
                    child_bounds[id(var1)] = existing.intersect(new_bound)
                else:
                    child_bounds[id(var1)] = Interval(new_bound.lower_bound, new_bound.upper_bound)

            else:
                new_bound_var1 = term_bound / bounds[var2]
                new_bound_var2 = term_bound / bounds[var1]

                if id(var1) in child_bounds:
                    existing = child_bounds[id(var1)]
                    child_bounds[id(var1)] = existing.intersect(new_bound_var1)
                else:
                    child_bounds[id(var1)] = new_bound_var1

                if id(var2) in child_bounds:
                    existing = child_bounds[id(var2)]
                    child_bounds[id(var2)] = existing.intersect(new_bound_var2)
                else:
                    child_bounds[id(var2)] = new_bound_var2

        return [child_bounds[id(v)] for v in expr.args]

    def _term_bound(self, term, bounds):
        if term.var1 is term.var2:
            return term.coefficient * (bounds[term.var1] ** 2)
        return term.coefficient * bounds[term.var1] * bounds[term.var2]


class PowerRule(Rule):
    """Return new bounds for power expressions."""
    def apply(self, expr, bounds):
        base, expo = expr.args
        if type(expo) not in nonpyomo_leaf_types:
            if not expo.is_constant():
                return None
            expo = expo.value

        if not almosteq(expo, 2):
            return None

        expr_bound = bounds[expr]
        # the bound of a square number is never negative, but check anyway to
        # avoid unexpected crashes.
        if not expr_bound.is_nonnegative():
            return None

        sqrt_bound = expr_bound.sqrt()
        return [
            Interval(-sqrt_bound.upper_bound, sqrt_bound.upper_bound),
            None,
        ]


class _UnaryFunctionRule(UnaryFunctionRule):
    def apply(self, expr, bounds):
        expr_bounds = bounds[expr]
        return [self._child_bounds(expr_bounds)]

    def _child_bounds(self, bounds):
        pass


class _BoundedFunctionRule(_UnaryFunctionRule):
    func_name = None

    def _child_bounds(self, bounds):
        func = getattr(bounds, self.func_name)
        return func.inverse()


class AbsRule(_BoundedFunctionRule):
    """Return new bounds for abs."""
    func_name = 'abs'


class SqrtRule(_BoundedFunctionRule):
    """Return new bounds for sqrt."""
    func_name = 'sqrt'


class ExpRule(_BoundedFunctionRule):
    """Return new bounds for exp."""
    func_name = 'exp'


class LogRule(_BoundedFunctionRule):
    """Return new bounds for log."""
    func_name = 'log'


class _UnboundedFunctionRule(_UnaryFunctionRule):
    def _child_bounds(self, bounds):
        return Interval(None, None)


class SinRule(_UnboundedFunctionRule):
    """Return new bounds for sin."""
    func_type = UnaryFunctionType.Sin


class CosRule(_UnboundedFunctionRule):
    """Return new bounds for cos."""
    func_type = UnaryFunctionType.Cos


class TanRule(_UnboundedFunctionRule):
    """Return new bounds for tan."""
    func_type = UnaryFunctionType.Tan


class AsinRule(_UnboundedFunctionRule):
    """Return new bounds for asin."""
    func_type = UnaryFunctionType.Asin


class AcosRule(_UnboundedFunctionRule):
    """Return new bounds for acos."""
    func_type = UnaryFunctionType.Acos


class AtanRule(_UnboundedFunctionRule):
    """Return new bounds for atan."""
    func_type = UnaryFunctionType.Atan


_func_name_to_rule_map = dict()
_func_name_to_rule_map['abs'] = AbsRule()
_func_name_to_rule_map['sqrt'] = SqrtRule()
_func_name_to_rule_map['exp'] = ExpRule()
_func_name_to_rule_map['log'] = LogRule()
_func_name_to_rule_map['sin'] = SinRule()
_func_name_to_rule_map['cos'] = CosRule()
_func_name_to_rule_map['tan'] = TanRule()
_func_name_to_rule_map['asin'] = AsinRule()
_func_name_to_rule_map['acos'] = AcosRule()
_func_name_to_rule_map['atan'] = AtanRule()


class UnaryFunctionRule(Rule):
    """Bound tightening rule for unary functions."""
    def apply(self, expr, bounds):
        assert len(expr.args) == 1
        func_name = expr.getname()
        if func_name not in _func_name_to_rule_map:
            raise ValueError('Unknown function type {}'.format(func_name))
        return _func_name_to_rule_map[func_name].apply(expr, bounds)
