"""
Enable optimizations on the math module.
"""

from astoptimizer.compatibility import INT_TYPES, FLOAT_TYPES
from astoptimizer.config import Function
import math
import sys

def check_sqrt(config, args):
    arg = args[0]
    return (0.0 <= arg)

def check_pow(config, args):
    num = args[0]
    exp = args[1]
    if exp < 1.0 and exp != 0.0 and num < 0:
        return False
    return True

def check_sum_args(config, args):
    data = args[0]
    return all(isinstance(item, FLOAT_TYPES) for item in data)

def setup_config(config):
    config.add_constant('math.e', math.e)
    config.add_constant('math.pi', math.pi)

    config.add_func('math.ceil', Function(math.ceil, 1, FLOAT_TYPES))
    config.add_func('math.exp', Function(math.exp, 1, FLOAT_TYPES))
    config.add_func('math.fabs', Function(math.fabs, 1, FLOAT_TYPES))
    config.add_func('math.floor', Function(math.floor, 1, FLOAT_TYPES))
    config.add_func('math.fmod', Function(math.fmod, 2, FLOAT_TYPES, FLOAT_TYPES))
    config.add_func('math.frexp', Function(math.frexp, 1, FLOAT_TYPES))
    config.add_func('math.ldexp', Function(math.ldexp, 2, FLOAT_TYPES, FLOAT_TYPES))
    config.add_func('math.modf', Function(math.modf, 1, FLOAT_TYPES))
    config.add_func('math.copysign', Function(math.copysign, 2, FLOAT_TYPES, FLOAT_TYPES))
    config.add_func('math.factorial', Function(math.factorial, 1, INT_TYPES))
    config.add_func('math.fsum', Function(math.fsum, 1, (tuple, frozenset), check_args=check_sum_args))
    config.add_func('math.log1p', Function(math.log1p, 1, FLOAT_TYPES))
    config.add_func('math.trunc', Function(math.trunc, 1, FLOAT_TYPES))
    if sys.version_info >= (3, 2):
        config.add_func('math.expm1', Function(math.expm1, 1, FLOAT_TYPES))
        config.add_func('math.isfinite', Function(math.isfinite, 1, FLOAT_TYPES))
        config.add_func('math.isinf', Function(math.isinf, 1, FLOAT_TYPES))
        config.add_func('math.isnan', Function(math.isnan, 1, FLOAT_TYPES))

    config.add_func('math.log', Function(math.log, (1, 2), FLOAT_TYPES, FLOAT_TYPES))
    if sys.version_info >= (3, 3):
        config.add_func('math.log2', Function(math.log2, 1, FLOAT_TYPES))
    config.add_func('math.log10', Function(math.log10, 1, FLOAT_TYPES))
    config.add_func('math.pow', Function(math.pow, 2, FLOAT_TYPES, FLOAT_TYPES, check_args=check_pow))
    config.add_func('math.sqrt', Function(math.sqrt, 1, FLOAT_TYPES, check_args=check_sqrt))

    # Trigonometric functions
    config.add_func('math.acos', Function(math.acos, 1, FLOAT_TYPES))
    config.add_func('math.asin', Function(math.asin, 1, FLOAT_TYPES))
    config.add_func('math.atan', Function(math.atan, 1, FLOAT_TYPES))
    config.add_func('math.atan2', Function(math.atan2, 2, FLOAT_TYPES, FLOAT_TYPES))
    config.add_func('math.cos', Function(math.cos, 1, FLOAT_TYPES))
    config.add_func('math.sin', Function(math.sin, 1, FLOAT_TYPES))
    config.add_func('math.tan', Function(math.tan, 1, FLOAT_TYPES))
    config.add_func('math.hypot', Function(math.hypot, 2, FLOAT_TYPES, FLOAT_TYPES))


