#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

from sympy import symbols, Eq, solve, StrictInequality, Interval

class Switch(object):
    """
    >>> BOUNDS = Switch("MIN", "MAX")
    >>> BOUNDS
    'MIN'
    >>> BOUNDS.MIN
    'MIN'
    >>> ~BOUNDS.MIN
    'MAX'
    >>> BOUNDS.MAX == ~BOUNDS.MIN
    True
    """
    def __init__(self, state1, state2):
        self.state1 = state1
        self.state2 = state2
        self.state = state1
    def __repr__(self):
        return repr(self.state)
    def __eq__(self, other_switch):
        return self.state == other_switch.state
    def __getattr__(self, state):
        if state == self.state1:
            return self.__class__(self.state1, self.state2)
        if state == self.state2:
            return self.__class__(self.state2, self.state1)
        return None
    def __invert__(self):
        if self.state == self.state2:
            return self.__class__(self.state1, self.state2)
        else:
            return self.__class__(self.state2, self.state1)

BOUNDS = Switch("MIN", "MAX")

def solve_inequality(equation, symbol):
    return solve(Eq(equation.lhs, equation.rhs), symbol)[0]

def symbol_bounds(equation, symbol, other_symbol):
    """
    >>> x, y = symbols("x y")
    >>> symbol_bounds(x > y, x, y)
    ('MIN', 'MIN')
    >>> symbol_bounds(x < y, x, y)
    ('MAX', 'MAX')
    >>> symbol_bounds(x + y > 0, x, y)
    ('MIN', 'MAX')
    >>> symbol_bounds(x + y < 0, x, y)
    ('MAX', 'MIN')
    """
    symbol_term_is_lower = symbol in equation.lhs
    other_symbol_term_is_lower = other_symbol in equation.lhs
    symbol_bounds = BOUNDS.MAX if symbol_term_is_lower else BOUNDS.MIN
    other_symbol_bounds = BOUNDS.MIN if other_symbol_term_is_lower \
        else BOUNDS.MAX
    return (symbol_bounds, other_symbol_bounds)

def bounds_offset(equation, symbol):
    """
    >>> x, y = symbols("x y")
    >>> bounds_offset(x >= 0, x)
    0
    >>> bounds_offset(x > 0, x)
    1
    >>> bounds_offset(x <= 0, x)
    0
    >>> bounds_offset(x < 0, x)
    -1
    """
    offset = 1 if isinstance(equation, StrictInequality) else 0
    return -offset if symbol in equation.lhs else offset

def get_bounds(interval, bounds):
    """
    >>> interval = Interval(4, 7)
    >>> get_bounds(interval, BOUNDS.MIN)
    4
    >>> get_bounds(interval, BOUNDS.MAX)
    7
    """
    if bounds == BOUNDS.MIN:
        return interval.start
    if bounds == BOUNDS.MAX:
        return interval.end
    return None

def set_bounds(interval, bounds, value):
    """
    >>> interval = Interval(4, 7)
    >>> set_bounds(interval, BOUNDS.MIN, 3)
    [3, 7]
    >>> set_bounds(interval, BOUNDS.MAX, 8)
    [4, 8]
    """
    if bounds == BOUNDS.MIN:
        return Interval(value, interval.end)
    if bounds == BOUNDS.MAX:
        return Interval(interval.start, value)
    return None

def bounds_consistency(constraint, symbol_1, symbol_2, domain_1, domain_2):
    """
    >>> x, y = symbols("x y")
    >>> domain_x = Interval(1, 10)
    >>> domain_y = Interval(1, 10)

    >>> bounds_consistency(x < y, x, y, domain_x, domain_y)
    [[1, 9], [2, 10]]
    >>> bounds_consistency(x + 2 < y, x, y, domain_x, domain_y)
    [[1, 7], [4, 10]]
    """
    def revise_domain(constraint, symbol, other_symbol, domain, other_domain):
        bounds, other_bounds = symbol_bounds(constraint, symbol, other_symbol)
        transposed_constraint = solve_inequality(constraint, symbol)
        offset = bounds_offset(constraint, symbol)
        substitution = get_bounds(other_domain, other_bounds) + offset
        result = transposed_constraint.subs(other_symbol, substitution)
        domain = set_bounds(domain, bounds, result)
        return domain
    domain_1 = revise_domain(constraint, symbol_1, symbol_2, domain_1, domain_2)
    domain_2 = revise_domain(constraint, symbol_2, symbol_1, domain_2, domain_1)
    return [domain_1, domain_2]

if __name__ == "__main__":
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    x, y = symbols("x y")
    domain_x = Interval(1, 10)
    domain_y = Interval(1, 10)

    domain_x, domain_y = bounds_consistency(x + 2 < y, x, y,
                                            domain_x, domain_y)
    print domain_x, domain_y

    domain_x, domain_y = bounds_consistency(x + y >= 16, x, y,
                                            domain_x, domain_y)
    print domain_x, domain_y