import sys
import platform
import pytest

import numpy as np
# import the c-extension module directly since _arg is not exported via umath
import numpy.core._multiarray_umath as ncu
from numpy.testing import (
    assert_raises, assert_equal, assert_array_equal, assert_almost_equal, assert_array_max_ulp
    )

# TODO: branch cuts (use Pauli code)
# TODO: conj 'symmetry'
# TODO: FPU exceptions

# At least on Windows the results of many complex functions are not conforming
# to the C99 standard. See ticket 1574.
# Ditto for Solaris (ticket 1642) and OS X on PowerPC.
#FIXME: this will probably change when we require full C99 campatibility
with np.errstate(all='ignore'):
    functions_seem_flaky = ((np.exp(complex(np.inf, 0)).imag != 0)
                            or (np.log(complex(np.NZERO, 0)).imag != np.pi))
# TODO: replace with a check on whether platform-provided C99 funcs are used
xfail_complex_tests = (not sys.platform.startswith('linux') or functions_seem_flaky)

# TODO This can be xfail when the generator functions are got rid of.
platform_skip = pytest.mark.skipif(xfail_complex_tests,
                                   reason="Inadequate C99 complex support")



class TestCexp:
    def test_simple(self):
        check = check_complex_value
        f = np.exp

        check(f, 1, 0, np.exp(1), 0, False)
        check(f, 0, 1, np.cos(1), np.sin(1), False)

        ref = np.exp(1) * complex(np.cos(1), np.sin(1))
        check(f, 1, 1, ref.real, ref.imag, False)

    @platform_skip
    def test_special_values(self):
        # C99: Section G 6.3.1

        check = check_complex_value
        f = np.exp

        # cexp(+-0 + 0i) is 1 + 0i
        check(f, np.PZERO, 0, 1, 0, False)
        check(f, np.NZERO, 0, 1, 0, False)

        # cexp(x + infi) is nan + nani for finite x and raises 'invalid' FPU
        # exception
        check(f,  1, np.inf, np.nan, np.nan)
        check(f, -1, np.inf, np.nan, np.nan)
        check(f,  0, np.inf, np.nan, np.nan)

        # cexp(inf + 0i) is inf + 0i
        check(f,  np.inf, 0, np.inf, 0)

        # cexp(-inf + yi) is +0 * (cos(y) + i sin(y)) for finite y
        check(f,  -np.inf, 1, np.PZERO, np.PZERO)
        check(f,  -np.inf, 0.75 * np.pi, np.NZERO, np.PZERO)

        # cexp(inf + yi) is +inf * (cos(y) + i sin(y)) for finite y
        check(f,  np.inf, 1, np.inf, np.inf)
        check(f,  np.inf, 0.75 * np.pi, -np.inf, np.inf)

        # cexp(-inf + inf i) is +-0 +- 0i (signs unspecified)
        def _check_ninf_inf(dummy):
            msgform = "cexp(-inf, inf) is (%f, %f), expected (+-0, +-0)"
            with np.errstate(invalid='ignore'):
                z = f(np.array(complex(-np.inf, np.inf)))
                if z.real != 0 or z.imag != 0:
                    raise AssertionError(msgform % (z.real, z.imag))

        _check_ninf_inf(None)

        # cexp(inf + inf i) is +-inf + NaNi and raised invalid FPU ex.
        def _check_inf_inf(dummy):
            msgform = "cexp(inf, inf) is (%f, %f), expected (+-inf, nan)"
            with np.errstate(invalid='ignore'):
                z = f(np.array(complex(np.inf, np.inf)))
                if not np.isinf(z.real) or not np.isnan(z.imag):
                    raise AssertionError(msgform % (z.real, z.imag))

        _check_inf_inf(None)

        # cexp(-inf + nan i) is +-0 +- 0i
        def _check_ninf_nan(dummy):
            msgform = "cexp(-inf, nan) is (%f, %f), expected (+-0, +-0)"
            with np.errstate(invalid='ignore'):
                z = f(np.array(complex(-np.inf, np.nan)))
                if z.real != 0 or z.imag != 0:
                    raise AssertionError(msgform % (z.real, z.imag))

        _check_ninf_nan(None)

        # cexp(inf + nan i) is +-inf + nan
        def _check_inf_nan(dummy):
            msgform = "cexp(-inf, nan) is (%f, %f), expected (+-inf, nan)"
            with np.errstate(invalid='ignore'):
                z = f(np.array(complex(np.inf, np.nan)))
                if not np.isinf(z.real) or not np.isnan(z.imag):
                    raise AssertionError(msgform % (z.real, z.imag))

        _check_inf_nan(None)

        # cexp(nan + yi) is nan + nani for y != 0 (optional: raises invalid FPU
        # ex)
        check(f, np.nan, 1, np.nan, np.nan)
        check(f, np.nan, -1, np.nan, np.nan)

        check(f, np.nan,  np.inf, np.nan, np.nan)
        check(f, np.nan, -np.inf, np.nan, np.nan)

        # cexp(nan + nani) is nan + nani
        check(f, np.nan, np.nan, np.nan, np.nan)

    # TODO This can be xfail when the generator functions are got rid of.
    @pytest.mark.skip(reason="cexp(nan + 0I) is wrong on most platforms")
    def test_special_values2(self):
        # XXX: most implementations get it wrong here (including glibc <= 2.10)
        # cexp(nan + 0i) is nan + 0i
        check = check_complex_value
        f = np.exp

        check(f, np.nan, 0, np.nan, 0)

class TestClog:
    def test_simple(self):
        x = np.array([1+0j, 1+2j])
        y_r = np.log(np.abs(x)) + 1j * np.angle(x)
        y = np.log(x)
        for i in range(len(x)):
            assert_almost_equal(y[i], y_r[i])

    @platform_skip
    @pytest.mark.skipif(platform.machine() == "armv5tel", reason="See gh-413.")
    def test_special_values(self):
        xl = []
        yl = []

        # From C99 std (Sec 6.3.2)
        # XXX: check exceptions raised
        # --- raise for invalid fails.

        # clog(-0 + i0) returns -inf + i pi and raises the 'divide-by-zero'
        # floating-point exception.
        with np.errstate(divide='raise'):
            x = np.array([np.NZERO], dtype=complex)
            y = complex(-np.inf, np.pi)
            assert_raises(FloatingPointError, np.log, x)
        with np.errstate(divide='ignore'):
            assert_almost_equal(np.log(x), y)

        xl.append(x)
        yl.append(y)

        # clog(+0 + i0) returns -inf + i0 and raises the 'divide-by-zero'
        # floating-point exception.
        with np.errstate(divide='raise'):
            x = np.array([0], dtype=complex)
            y = complex(-np.inf, 0)
            assert_raises(FloatingPointError, np.log, x)
        with np.errstate(divide='ignore'):
            assert_almost_equal(np.log(x), y)

        xl.append(x)
        yl.append(y)

        # clog(x + i inf returns +inf + i pi /2, for finite x.
        x = np.array([complex(1, np.inf)], dtype=complex)
        y = complex(np.inf, 0.5 * np.pi)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        x = np.array([complex(-1, np.inf)], dtype=complex)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(x + iNaN) returns NaN + iNaN and optionally raises the
        # 'invalid' floating- point exception, for finite x.
        with np.errstate(invalid='raise'):
            x = np.array([complex(1., np.nan)], dtype=complex)
            y = complex(np.nan, np.nan)
            #assert_raises(FloatingPointError, np.log, x)
        with np.errstate(invalid='ignore'):
            assert_almost_equal(np.log(x), y)

        xl.append(x)
        yl.append(y)

        with np.errstate(invalid='raise'):
            x = np.array([np.inf + 1j * np.nan], dtype=complex)
            #assert_raises(FloatingPointError, np.log, x)
        with np.errstate(invalid='ignore'):
            assert_almost_equal(np.log(x), y)

        xl.append(x)
        yl.append(y)

        # clog(- inf + iy) returns +inf + ipi , for finite positive-signed y.
        x = np.array([-np.inf + 1j], dtype=complex)
        y = complex(np.inf, np.pi)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(+ inf + iy) returns +inf + i0, for finite positive-signed y.
        x = np.array([np.inf + 1j], dtype=complex)
        y = complex(np.inf, 0)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(- inf + i inf) returns +inf + i3pi /4.
        x = np.array([complex(-np.inf, np.inf)], dtype=complex)
        y = complex(np.inf, 0.75 * np.pi)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(+ inf + i inf) returns +inf + ipi /4.
        x = np.array([complex(np.inf, np.inf)], dtype=complex)
        y = complex(np.inf, 0.25 * np.pi)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(+/- inf + iNaN) returns +inf + iNaN.
        x = np.array([complex(np.inf, np.nan)], dtype=complex)
        y = complex(np.inf, np.nan)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        x = np.array([complex(-np.inf, np.nan)], dtype=complex)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(NaN + iy) returns NaN + iNaN and optionally raises the
        # 'invalid' floating-point exception, for finite y.
        x = np.array([complex(np.nan, 1)], dtype=complex)
        y = complex(np.nan, np.nan)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(NaN + i inf) returns +inf + iNaN.
        x = np.array([complex(np.nan, np.inf)], dtype=complex)
        y = complex(np.inf, np.nan)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(NaN + iNaN) returns NaN + iNaN.
        x = np.array([complex(np.nan, np.nan)], dtype=complex)
        y = complex(np.nan, np.nan)
        assert_almost_equal(np.log(x), y)
        xl.append(x)
        yl.append(y)

        # clog(conj(z)) = conj(clog(z)).
        xa = np.array(xl, dtype=complex)
        ya = np.array(yl, dtype=complex)
        with np.errstate(divide='ignore'):
            for i in range(len(xa)):
                assert_almost_equal(np.log(xa[i].conj()), ya[i].conj())


class TestCsqrt:

    def test_simple(self):
        # sqrt(1)
        check_complex_value(np.sqrt, 1, 0, 1, 0)

        # sqrt(1i)
        rres = 0.5*np.sqrt(2)
        ires = rres
        check_complex_value(np.sqrt, 0, 1, rres, ires, False)

        # sqrt(-1)
        check_complex_value(np.sqrt, -1, 0, 0, 1)

    def test_simple_conjugate(self):
        ref = np.conj(np.sqrt(complex(1, 1)))

        def f(z):
            return np.sqrt(np.conj(z))

        check_complex_value(f, 1, 1, ref.real, ref.imag, False)

    #def test_branch_cut(self):
    #    _check_branch_cut(f, -1, 0, 1, -1)

    @platform_skip
    def test_special_values(self):
        # C99: Sec G 6.4.2

        check = check_complex_value
        f = np.sqrt

        # csqrt(+-0 + 0i) is 0 + 0i
        check(f, np.PZERO, 0, 0, 0)
        check(f, np.NZERO, 0, 0, 0)

        # csqrt(x + infi) is inf + infi for any x (including NaN)
        check(f,  1, np.inf, np.inf, np.inf)
        check(f, -1, np.inf, np.inf, np.inf)

        check(f, np.PZERO, np.inf, np.inf, np.inf)
        check(f, np.NZERO, np.inf, np.inf, np.inf)
        check(f,   np.inf, np.inf, np.inf, np.inf)
        check(f,  -np.inf, np.inf, np.inf, np.inf)
        check(f,  -np.nan, np.inf, np.inf, np.inf)

        # csqrt(x + nani) is nan + nani for any finite x
        check(f,  1, np.nan, np.nan, np.nan)
        check(f, -1, np.nan, np.nan, np.nan)
        check(f,  0, np.nan, np.nan, np.nan)

        # csqrt(-inf + yi) is +0 + infi for any finite y > 0
        check(f, -np.inf, 1, np.PZERO, np.inf)

        # csqrt(inf + yi) is +inf + 0i for any finite y > 0
        check(f, np.inf, 1, np.inf, np.PZERO)

        # csqrt(-inf + nani) is nan +- infi (both +i infi are valid)
        def _check_ninf_nan(dummy):
            msgform = "csqrt(-inf, nan) is (%f, %f), expected (nan, +-inf)"
            z = np.sqrt(np.array(complex(-np.inf, np.nan)))
            #Fixme: ugly workaround for isinf bug.
            with np.errstate(invalid='ignore'):
                if not (np.isnan(z.real) and np.isinf(z.imag)):
                    raise AssertionError(msgform % (z.real, z.imag))

        _check_ninf_nan(None)

        # csqrt(+inf + nani) is inf + nani