1"""
2Low-level functions for complex arithmetic.
3"""
4
5import sys
6
7from .backend import MPZ, MPZ_ZERO, MPZ_ONE, MPZ_TWO, BACKEND
8
9from .libmpf import (\
10    round_floor, round_ceiling, round_down, round_up,
11    round_nearest, round_fast, bitcount,
12    bctable, normalize, normalize1, reciprocal_rnd, rshift, lshift, giant_steps,
13    negative_rnd,
14    to_str, to_fixed, from_man_exp, from_float, to_float, from_int, to_int,
15    fzero, fone, ftwo, fhalf, finf, fninf, fnan, fnone,
16    mpf_abs, mpf_pos, mpf_neg, mpf_add, mpf_sub, mpf_mul,
17    mpf_div, mpf_mul_int, mpf_shift, mpf_sqrt, mpf_hypot,
18    mpf_rdiv_int, mpf_floor, mpf_ceil, mpf_nint, mpf_frac,
19    mpf_sign, mpf_hash,
20    ComplexResult
21)
22
23from .libelefun import (\
24    mpf_pi, mpf_exp, mpf_log, mpf_cos_sin, mpf_cosh_sinh, mpf_tan, mpf_pow_int,
25    mpf_log_hypot,
26    mpf_cos_sin_pi, mpf_phi,
27    mpf_cos, mpf_sin, mpf_cos_pi, mpf_sin_pi,
28    mpf_atan, mpf_atan2, mpf_cosh, mpf_sinh, mpf_tanh,
29    mpf_asin, mpf_acos, mpf_acosh, mpf_nthroot, mpf_fibonacci
30)
31
32# An mpc value is a (real, imag) tuple
33mpc_one = fone, fzero
34mpc_zero = fzero, fzero
35mpc_two = ftwo, fzero
36mpc_half = (fhalf, fzero)
37
38_infs = (finf, fninf)
39_infs_nan = (finf, fninf, fnan)
40
41def mpc_is_inf(z):
42    """Check if either real or imaginary part is infinite"""
43    re, im = z
44    if re in _infs: return True
45    if im in _infs: return True
46    return False
47
48def mpc_is_infnan(z):
49    """Check if either real or imaginary part is infinite or nan"""
50    re, im = z
51    if re in _infs_nan: return True
52    if im in _infs_nan: return True
53    return False
54
55def mpc_to_str(z, dps, **kwargs):
56    re, im = z
57    rs = to_str(re, dps)
58    if im[0]:
59        return rs + " - " + to_str(mpf_neg(im), dps, **kwargs) + "j"
60    else:
61        return rs + " + " + to_str(im, dps, **kwargs) + "j"
62
63def mpc_to_complex(z, strict=False, rnd=round_fast):
64    re, im = z
65    return complex(to_float(re, strict, rnd), to_float(im, strict, rnd))
66
67def mpc_hash(z):
68    if sys.version_info >= (3, 2):
69        re, im = z
70        h = mpf_hash(re) + sys.hash_info.imag * mpf_hash(im)
71        # Need to reduce either module 2^32 or 2^64
72        h = h % (2**sys.hash_info.width)
73        return int(h)
74    else:
75        try:
76            return hash(mpc_to_complex(z, strict=True))
77        except OverflowError:
78            return hash(z)
79
80def mpc_conjugate(z, prec, rnd=round_fast):
81    re, im = z
82    return re, mpf_neg(im, prec, rnd)
83
84def mpc_is_nonzero(z):
85    return z != mpc_zero
86
87def mpc_add(z, w, prec, rnd=round_fast):
88    a, b = z
89    c, d = w
90    return mpf_add(a, c, prec, rnd), mpf_add(b, d, prec, rnd)
91
92def mpc_add_mpf(z, x, prec, rnd=round_fast):
93    a, b = z
94    return mpf_add(a, x, prec, rnd), b
95
96def mpc_sub(z, w, prec=0, rnd=round_fast):
97    a, b = z
98    c, d = w
99    return mpf_sub(a, c, prec, rnd), mpf_sub(b, d, prec, rnd)
100
101def mpc_sub_mpf(z, p, prec=0, rnd=round_fast):
102    a, b = z
103    return mpf_sub(a, p, prec, rnd), b
104
105def mpc_pos(z, prec, rnd=round_fast):
106    a, b = z
107    return mpf_pos(a, prec, rnd), mpf_pos(b, prec, rnd)
108
109def mpc_neg(z, prec=None, rnd=round_fast):
110    a, b = z
111    return mpf_neg(a, prec, rnd), mpf_neg(b, prec, rnd)
112
113def mpc_shift(z, n):
114    a, b = z
115    return mpf_shift(a, n), mpf_shift(b, n)
116
117def mpc_abs(z, prec, rnd=round_fast):
118    """Absolute value of a complex number, |a+bi|.
119    Returns an mpf value."""
120    a, b = z
121    return mpf_hypot(a, b, prec, rnd)
122
123def mpc_arg(z, prec, rnd=round_fast):
124    """Argument of a complex number. Returns an mpf value."""
125    a, b = z
126    return mpf_atan2(b, a, prec, rnd)
127
128def mpc_floor(z, prec, rnd=round_fast):
129    a, b = z
130    return mpf_floor(a, prec, rnd), mpf_floor(b, prec, rnd)
131
132def mpc_ceil(z, prec, rnd=round_fast):
133    a, b = z
134    return mpf_ceil(a, prec, rnd), mpf_ceil(b, prec, rnd)
135
136def mpc_nint(z, prec, rnd=round_fast):
137    a, b = z
138    return mpf_nint(a, prec, rnd), mpf_nint(b, prec, rnd)
139
140def mpc_frac(z, prec, rnd=round_fast):
141    a, b = z
142    return mpf_frac(a, prec, rnd), mpf_frac(b, prec, rnd)
143
144
145def mpc_mul(z, w, prec, rnd=round_fast):
146    """
147    Complex multiplication.
148
149    Returns the real and imaginary part of (a+bi)*(c+di), rounded to
150    the specified precision. The rounding mode applies to the real and
151    imaginary parts separately.
152    """
153    a, b = z
154    c, d = w
155    p = mpf_mul(a, c)
156    q = mpf_mul(b, d)
157    r = mpf_mul(a, d)
158    s = mpf_mul(b, c)
159    re = mpf_sub(p, q, prec, rnd)
160    im = mpf_add(r, s, prec, rnd)
161    return re, im
162
163def mpc_square(z, prec, rnd=round_fast):
164    # (a+b*I)**2 == a**2 - b**2 + 2*I*a*b
165    a, b = z
166    p = mpf_mul(a,a)
167    q = mpf_mul(b,b)
168    r = mpf_mul(a,b, prec, rnd)
169    re = mpf_sub(p, q, prec, rnd)
170    im = mpf_shift(r, 1)
171    return re, im
172
173def mpc_mul_mpf(z, p, prec, rnd=round_fast):
174    a, b = z
175    re = mpf_mul(a, p, prec, rnd)
176    im = mpf_mul(b, p, prec, rnd)
177    return re, im
178
179def mpc_mul_imag_mpf(z, x, prec, rnd=round_fast):
180    """
181    Multiply the mpc value z by I*x where x is an mpf value.
182    """
183    a, b = z
184    re = mpf_neg(mpf_mul(b, x, prec, rnd))
185    im = mpf_mul(a, x, prec, rnd)
186    return re, im
187
188def mpc_mul_int(z, n, prec, rnd=round_fast):
189    a, b = z
190    re = mpf_mul_int(a, n, prec, rnd)
191    im = mpf_mul_int(b, n, prec, rnd)
192    return re, im
193
194def mpc_div(z, w, prec, rnd=round_fast):
195    a, b = z
196    c, d = w
197    wp = prec + 10
198    # mag = c*c + d*d
199    mag = mpf_add(mpf_mul(c, c), mpf_mul(d, d), wp)
200    # (a*c+b*d)/mag, (b*c-a*d)/mag
201    t = mpf_add(mpf_mul(a,c), mpf_mul(b,d), wp)
202    u = mpf_sub(mpf_mul(b,c), mpf_mul(a,d), wp)
203    return mpf_div(t,mag,prec,rnd), mpf_div(u,mag,prec,rnd)
204
205def mpc_div_mpf(z, p, prec, rnd=round_fast):
206    """Calculate z/p where p is real"""
207    a, b = z
208    re = mpf_div(a, p, prec, rnd)
209    im = mpf_div(b, p, prec, rnd)
210    return re, im
211
212def mpc_reciprocal(z, prec, rnd=round_fast):
213    """Calculate 1/z efficiently"""
214    a, b = z
215    m = mpf_add(mpf_mul(a,a),mpf_mul(b,b),prec+10)
216    re = mpf_div(a, m, prec, rnd)
217    im = mpf_neg(mpf_div(b, m, prec, rnd))
218    return re, im
219
220def mpc_mpf_div(p, z, prec, rnd=round_fast):
221    """Calculate p/z where p is real efficiently"""
222    a, b = z
223    m = mpf_add(mpf_mul(a,a),mpf_mul(b,b), prec+10)
224    re = mpf_div(mpf_mul(a,p), m, prec, rnd)
225    im = mpf_div(mpf_neg(mpf_mul(b,p)), m, prec, rnd)
226    return re, im
227
228def complex_int_pow(a, b, n):
229    """Complex integer power: computes (a+b*I)**n exactly for
230    nonnegative n (a and b must be Python ints)."""
231    wre = 1
232    wim = 0
233    while n:
234        if n & 1:
235            wre, wim = wre*a - wim*b, wim*a + wre*b
236            n -= 1
237        a, b = a*a - b*b, 2*a*b
238        n //= 2
239    return wre, wim
240
241def mpc_pow(z, w, prec, rnd=round_fast):
242    if w[1] == fzero:
243        return mpc_pow_mpf(z, w[0], prec, rnd)
244    return mpc_exp(mpc_mul(mpc_log(z, prec+10), w, prec+10), prec, rnd)
245
246def mpc_pow_mpf(z, p, prec, rnd=round_fast):
247    psign, pman, pexp, pbc = p
248    if pexp >= 0:
249        return mpc_pow_int(z, (-1)**psign * (pman<<pexp), prec, rnd)
250    if pexp == -1:
251        sqrtz = mpc_sqrt(z, prec+10)
252        return mpc_pow_int(sqrtz, (-1)**psign * pman, prec, rnd)
253    return mpc_exp(mpc_mul_mpf(mpc_log(z, prec+10), p, prec+10), prec, rnd)
254
255def mpc_pow_int(z, n, prec, rnd=round_fast):
256    a, b = z
257    if b == fzero:
258        return mpf_pow_int(a, n, prec, rnd), fzero
259    if a == fzero:
260        v = mpf_pow_int(b, n, prec, rnd)
261        n %= 4
262        if n == 0:
263            return v, fzero
264        elif n == 1:
265            return fzero, v
266        elif n == 2:
267            return mpf_neg(v), fzero
268        elif n == 3:
269            return fzero, mpf_neg(v)
270    if n == 0: return mpc_one
271    if n == 1: return mpc_pos(z, prec, rnd)
272    if n == 2: return mpc_square(z, prec, rnd)
273    if n == -1: return mpc_reciprocal(z, prec, rnd)
274    if n < 0: return mpc_reciprocal(mpc_pow_int(z, -n, prec+4), prec, rnd)
275    asign, aman, aexp, abc = a
276    bsign, bman, bexp, bbc = b
277    if asign: aman = -aman
278    if bsign: bman = -bman
279    de = aexp - bexp
280    abs_de = abs(de)
281    exact_size = n*(abs_de + max(abc, bbc))
282    if exact_size < 10000:
283        if de > 0:
284            aman <<= de
285            aexp = bexp
286        else:
287            bman <<= (-de)
288            bexp = aexp
289        re, im = complex_int_pow(aman, bman, n)
290        re = from_man_exp(re, int(n*aexp), prec, rnd)
291        im = from_man_exp(im, int(n*bexp), prec, rnd)
292        return re, im
293    return mpc_exp(mpc_mul_int(mpc_log(z, prec+10), n, prec+10), prec, rnd)
294
295def mpc_sqrt(z, prec, rnd=round_fast):
296    """Complex square root (principal branch).
297
298    We have sqrt(a+bi) = sqrt((r+a)/2) + b/sqrt(2*(r+a))*i where
299    r = abs(a+bi), when a+bi is not a negative real number."""
300    a, b = z
301    if b == fzero:
302        if a == fzero:
303            return (a, b)
304        # When a+bi is a negative real number, we get a real sqrt times i
305        if a[0]:
306            im = mpf_sqrt(mpf_neg(a), prec, rnd)
307            return (fzero, im)
308        else:
309            re = mpf_sqrt(a, prec, rnd)
310            return (re, fzero)
311    wp = prec+20
312    if not a[0]:                               # case a positive
313        t  = mpf_add(mpc_abs((a, b), wp), a, wp)  # t = abs(a+bi) + a
314        u = mpf_shift(t, -1)                      # u = t/2
315        re = mpf_sqrt(u, prec, rnd)               # re = sqrt(u)
316        v = mpf_shift(t, 1)                       # v = 2*t
317        w  = mpf_sqrt(v, wp)                      # w = sqrt(v)
318        im = mpf_div(b, w, prec, rnd)             # im = b / w
319    else:                                      # case a negative
320        t = mpf_sub(mpc_abs((a, b), wp), a, wp)   # t = abs(a+bi) - a
321        u = mpf_shift(t, -1)                      # u = t/2
322        im = mpf_sqrt(u, prec, rnd)               # im = sqrt(u)
323        v = mpf_shift(t, 1)                       # v = 2*t
324        w  = mpf_sqrt(v, wp)                      # w = sqrt(v)
325        re = mpf_div(b, w, prec, rnd)             # re = b/w
326        if b[0]:
327            re = mpf_neg(re)
328            im = mpf_neg(im)
329    return re, im
330
331def mpc_nthroot_fixed(a, b, n, prec):
332    # a, b signed integers at fixed precision prec
333    start = 50
334    a1 = int(rshift(a, prec - n*start))
335    b1 = int(rshift(b, prec - n*start))
336    try:
337        r = (a1 + 1j * b1)**(1.0/n)
338        re = r.real
339        im = r.imag
340        re = MPZ(int(re))
341        im = MPZ(int(im))
342    except OverflowError:
343        a1 = from_int(a1, start)
344        b1 = from_int(b1, start)
345        fn = from_int(n)
346        nth = mpf_rdiv_int(1, fn, start)
347        re, im = mpc_pow((a1, b1), (nth, fzero), start)
348        re = to_int(re)
349        im = to_int(im)
350    extra = 10
351    prevp = start
352    extra1 = n
353    for p in giant_steps(start, prec+extra):
354        # this is slow for large n, unlike int_pow_fixed
355        re2, im2 = complex_int_pow(re, im, n-1)
356        re2 = rshift(re2, (n-1)*prevp - p - extra1)
357        im2 = rshift(im2, (n-1)*prevp - p - extra1)
358        r4 = (re2*re2 + im2*im2) >> (p + extra1)
359        ap = rshift(a, prec - p)
360        bp = rshift(b, prec - p)
361        rec = (ap * re2 + bp * im2) >> p
362        imc = (-ap * im2 + bp * re2) >> p
363        reb = (rec << p) // r4
364        imb = (imc << p) // r4
365        re = (reb + (n-1)*lshift(re, p-prevp))//n
366        im = (imb + (n-1)*lshift(im, p-prevp))//n
367        prevp = p
368    return re, im
369
370def mpc_nthroot(z, n, prec, rnd=round_fast):
371    """
372    Complex n-th root.
373
374    Use Newton method as in the real case when it is faster,
375    otherwise use z**(1/n)
376    """
377    a, b = z
378    if a[0] == 0 and b == fzero:
379        re = mpf_nthroot(a, n, prec, rnd)
380        return (re, fzero)
381    if n < 2:
382        if n == 0:
383            return mpc_one
384        if n == 1:
385            return mpc_pos((a, b), prec, rnd)
386        if n == -1:
387            return mpc_div(mpc_one, (a, b), prec, rnd)
388        inverse = mpc_nthroot((a, b), -n, prec+5, reciprocal_rnd[rnd])
389        return mpc_div(mpc_one, inverse, prec, rnd)
390    if n <= 20:
391        prec2 = int(1.2 * (prec + 10))
392        asign, aman, aexp, abc = a
393        bsign, bman, bexp, bbc = b
394        pf = mpc_abs((a,b), prec)
395        if pf[-2] + pf[-1] > -10  and pf[-2] + pf[-1] < prec:
396            af = to_fixed(a, prec2)
397            bf = to_fixed(b, prec2)
398            re, im = mpc_nthroot_fixed(af, bf, n, prec2)
399            extra = 10
400            re = from_man_exp(re, -prec2-extra, prec2, rnd)
401            im = from_man_exp(im, -prec2-extra, prec2, rnd)
402            return re, im
403    fn = from_int(n)
404    prec2 = prec+10 + 10
405    nth = mpf_rdiv_int(1, fn, prec2)
406    re, im = mpc_pow((a, b), (nth, fzero), prec2, rnd)
407    re = normalize(re[0], re[1], re[2], re[3], prec, rnd)
408    im = normalize(im[0], im[1], im[2], im[3], prec, rnd)
409    return re, im
410
411def mpc_cbrt(z, prec, rnd=round_fast):
412    """
413    Complex cubic root.
414    """
415    return mpc_nthroot(z, 3, prec, rnd)
416
417def mpc_exp(z, prec, rnd=round_fast):
418    """
419    Complex exponential function.
420
421    We use the direct formula exp(a+bi) = exp(a) * (cos(b) + sin(b)*i)
422    for the computation. This formula is very nice because it is
423    pefectly stable; since we just do real multiplications, the only
424    numerical errors that can creep in are single-ulp rounding errors.
425
426    The formula is efficient since mpmath's real exp is quite fast and
427    since we can compute cos and sin simultaneously.
428
429    It is no problem if a and b are large; if the implementations of
430    exp/cos/sin are accurate and efficient for all real numbers, then
431    so is this function for all complex numbers.
432    """
433    a, b = z
434    if a == fzero:
435        return mpf_cos_sin(b, prec, rnd)
436    if b == fzero:
437        return mpf_exp(a, prec, rnd), fzero
438    mag = mpf_exp(a, prec+4, rnd)
439    c, s = mpf_cos_sin(b, prec+4, rnd)
440    re = mpf_mul(mag, c, prec, rnd)
441    im = mpf_mul(mag, s, prec, rnd)
442    return re, im
443
444def mpc_log(z, prec, rnd=round_fast):
445    re = mpf_log_hypot(z[0], z[1], prec, rnd)
446    im = mpc_arg(z, prec, rnd)
447    return re, im
448
449def mpc_cos(z, prec, rnd=round_fast):
450    """Complex cosine. The formula used is cos(a+bi) = cos(a)*cosh(b) -
451    sin(a)*sinh(b)*i.
452
453    The same comments apply as for the complex exp: only real
454    multiplications are pewrormed, so no cancellation errors are
455    possible. The formula is also efficient since we can compute both
456    pairs (cos, sin) and (cosh, sinh) in single stwps."""
457    a, b = z
458    if b == fzero:
459        return mpf_cos(a, prec, rnd), fzero
460    if a == fzero:
461        return mpf_cosh(b, prec, rnd), fzero
462    wp = prec + 6
463    c, s = mpf_cos_sin(a, wp)
464    ch, sh = mpf_cosh_sinh(b, wp)
465    re = mpf_mul(c, ch, prec, rnd)
466    im = mpf_mul(s, sh, prec, rnd)
467    return re, mpf_neg(im)
468
469def mpc_sin(z, prec, rnd=round_fast):
470    """Complex sine. We have sin(a+bi) = sin(a)*cosh(b) +
471    cos(a)*sinh(b)*i. See the docstring for mpc_cos for additional
472    comments."""
473    a, b = z
474    if b == fzero:
475        return mpf_sin(a, prec, rnd), fzero
476    if a == fzero:
477        return fzero, mpf_sinh(b, prec, rnd)
478    wp = prec + 6
479    c, s = mpf_cos_sin(a, wp)
480    ch, sh = mpf_cosh_sinh(b, wp)
481    re = mpf_mul(s, ch, prec, rnd)
482    im = mpf_mul(c, sh, prec, rnd)
483    return re, im
484
485def mpc_tan(z, prec, rnd=round_fast):
486    """Complex tangent. Computed as tan(a+bi) = sin(2a)/M + sinh(2b)/M*i
487    where M = cos(2a) + cosh(2b)."""
488    a, b = z
489    asign, aman, aexp, abc = a
490    bsign, bman, bexp, bbc = b
491    if b == fzero: return mpf_tan(a, prec, rnd), fzero
492    if a == fzero: return fzero, mpf_tanh(b, prec, rnd)
493    wp = prec + 15
494    a = mpf_shift(a, 1)
495    b = mpf_shift(b, 1)
496    c, s = mpf_cos_sin(a, wp)
497    ch, sh = mpf_cosh_sinh(b, wp)
498    # TODO: handle cancellation when c ~=  -1 and ch ~= 1
499    mag = mpf_add(c, ch, wp)
500    re = mpf_div(s, mag, prec, rnd)
501    im = mpf_div(sh, mag, prec, rnd)
502    return re, im
503
504def mpc_cos_pi(z, prec, rnd=round_fast):
505    a, b = z
506    if b == fzero:
507        return mpf_cos_pi(a, prec, rnd), fzero
508    b = mpf_mul(b, mpf_pi(prec+5), prec+5)
509    if a == fzero:
510        return mpf_cosh(b, prec, rnd), fzero
511    wp = prec + 6
512    c, s = mpf_cos_sin_pi(a, wp)
513    ch, sh = mpf_cosh_sinh(b, wp)
514    re = mpf_mul(c, ch, prec, rnd)
515    im = mpf_mul(s, sh, prec, rnd)
516    return re, mpf_neg(im)
517
518def mpc_sin_pi(z, prec, rnd=round_fast):
519    a, b = z
520    if b == fzero:
521        return mpf_sin_pi(a, prec, rnd), fzero
522    b = mpf_mul(b, mpf_pi(prec+5), prec+5)
523    if a == fzero:
524        return fzero, mpf_sinh(b, prec, rnd)
525    wp = prec + 6
526    c, s = mpf_cos_sin_pi(a, wp)
527    ch, sh = mpf_cosh_sinh(b, wp)
528    re = mpf_mul(s, ch, prec, rnd)
529    im = mpf_mul(c, sh, prec, rnd)
530    return re, im
531
532def mpc_cos_sin(z, prec, rnd=round_fast):
533    a, b = z
534    if a == fzero:
535        ch, sh = mpf_cosh_sinh(b, prec, rnd)
536        return (ch, fzero), (fzero, sh)
537    if b == fzero:
538        c, s = mpf_cos_sin(a, prec, rnd)
539        return (c, fzero), (s, fzero)
540    wp = prec + 6
541    c, s = mpf_cos_sin(a, wp)
542    ch, sh = mpf_cosh_sinh(b, wp)
543    cre = mpf_mul(c, ch, prec, rnd)
544    cim = mpf_mul(s, sh, prec, rnd)
545    sre = mpf_mul(s, ch, prec, rnd)
546    sim = mpf_mul(c, sh, prec, rnd)
547    return (cre, mpf_neg(cim)), (sre, sim)
548
549def mpc_cos_sin_pi(z, prec, rnd=round_fast):
550    a, b = z
551    if b == fzero:
552        c, s = mpf_cos_sin_pi(a, prec, rnd)
553        return (c, fzero), (s, fzero)
554    b = mpf_mul(b, mpf_pi(prec+5), prec+5)
555    if a == fzero:
556        ch, sh = mpf_cosh_sinh(b, prec, rnd)
557        return (ch, fzero), (fzero, sh)
558    wp = prec + 6
559    c, s = mpf_cos_sin_pi(a, wp)
560    ch, sh = mpf_cosh_sinh(b, wp)
561    cre = mpf_mul(c, ch, prec, rnd)
562    cim = mpf_mul(s, sh, prec, rnd)
563    sre = mpf_mul(s, ch, prec, rnd)
564    sim = mpf_mul(c, sh, prec, rnd)
565    return (cre, mpf_neg(cim)), (sre, sim)
566
567def mpc_cosh(z, prec, rnd=round_fast):
568    """Complex hyperbolic cosine. Computed as cosh(z) = cos(z*i)."""
569    a, b = z
570    return mpc_cos((b, mpf_neg(a)), prec, rnd)
571
572def mpc_sinh(z, prec, rnd=round_fast):
573    """Complex hyperbolic sine. Computed as sinh(z) = -i*sin(z*i)."""
574    a, b = z
575    b, a = mpc_sin((b, a), prec, rnd)
576    return a, b
577
578def mpc_tanh(z, prec, rnd=round_fast):
579    """Complex hyperbolic tangent. Computed as tanh(z) = -i*tan(z*i)."""
580    a, b = z
581    b, a = mpc_tan((b, a), prec, rnd)
582    return a, b
583
584# TODO: avoid loss of accuracy
585def mpc_atan(z, prec, rnd=round_fast):
586    a, b = z
587    # atan(z) = (I/2)*(log(1-I*z) - log(1+I*z))
588    # x = 1-I*z = 1 + b - I*a
589    # y = 1+I*z = 1 - b + I*a
590    wp = prec + 15
591    x = mpf_add(fone, b, wp), mpf_neg(a)
592    y = mpf_sub(fone, b, wp), a
593    l1 = mpc_log(x, wp)
594    l2 = mpc_log(y, wp)
595    a, b = mpc_sub(l1, l2, prec, rnd)
596    # (I/2) * (a+b*I) = (-b/2 + a/2*I)
597    v = mpf_neg(mpf_shift(b,-1)), mpf_shift(a,-1)
598    # Subtraction at infinity gives correct real part but
599    # wrong imaginary part (should be zero)
600    if v[1] == fnan and mpc_is_inf(z):
601        v = (v[0], fzero)
602    return v
603
604beta_crossover = from_float(0.6417)
605alpha_crossover = from_float(1.5)
606
607def acos_asin(z, prec, rnd, n):
608    """ complex acos for n = 0, asin for n = 1
609    The algorithm is described in
610    T.E. Hull, T.F. Fairgrieve and P.T.P. Tang
611    'Implementing the Complex Arcsine and Arcosine Functions
612    using Exception Handling',
613    ACM Trans. on Math. Software Vol. 23 (1997), p299
614    The complex acos and asin can be defined as
615    acos(z) = acos(beta) - I*sign(a)* log(alpha + sqrt(alpha**2 -1))
616    asin(z) = asin(beta) + I*sign(a)* log(alpha + sqrt(alpha**2 -1))
617    where z = a + I*b
618    alpha = (1/2)*(r + s); beta = (1/2)*(r - s) = a/alpha
619    r = sqrt((a+1)**2 + y**2); s = sqrt((a-1)**2 + y**2)
620    These expressions are rewritten in different ways in different
621    regions, delimited by two crossovers alpha_crossover and beta_crossover,
622    and by abs(a) <= 1, in order to improve the numerical accuracy.
623    """
624    a, b = z
625    wp = prec + 10
626    # special cases with real argument
627    if b == fzero:
628        am = mpf_sub(fone, mpf_abs(a), wp)
629        # case abs(a) <= 1
630        if not am[0]:
631            if n == 0:
632                return mpf_acos(a, prec, rnd), fzero
633            else:
634                return mpf_asin(a, prec, rnd), fzero
635        # cases abs(a) > 1
636        else:
637            # case a < -1
638            if a[0]:
639                pi = mpf_pi(prec, rnd)
640                c = mpf_acosh(mpf_neg(a), prec, rnd)
641                if n == 0:
642                    return pi, mpf_neg(c)
643                else:
644                    return mpf_neg(mpf_shift(pi, -1)), c
645            # case a > 1
646            else:
647                c = mpf_acosh(a, prec, rnd)
648                if n == 0:
649                    return fzero, c
650                else:
651                    pi = mpf_pi(prec, rnd)
652                    return mpf_shift(pi, -1), mpf_neg(c)
653    asign = bsign = 0
654    if a[0]:
655        a = mpf_neg(a)
656        asign = 1
657    if b[0]:
658        b = mpf_neg(b)
659        bsign = 1
660    am = mpf_sub(fone, a, wp)
661    ap = mpf_add(fone, a, wp)
662    r = mpf_hypot(ap, b, wp)
663    s = mpf_hypot(am, b, wp)
664    alpha = mpf_shift(mpf_add(r, s, wp), -1)
665    beta = mpf_div(a, alpha, wp)
666    b2 = mpf_mul(b,b, wp)
667    # case beta <= beta_crossover
668    if not mpf_sub(beta_crossover, beta, wp)[0]:
669        if n == 0:
670            re = mpf_acos(beta, wp)
671        else:
672            re = mpf_asin(beta, wp)
673    else:
674        # to compute the real part in this region use the identity
675        # asin(beta) = atan(beta/sqrt(1-beta**2))
676        # beta/sqrt(1-beta**2) = (alpha + a) * (alpha - a)
677        # alpha + a is numerically accurate; alpha - a can have
678        # cancellations leading to numerical inaccuracies, so rewrite
679        # it in differente ways according to the region
680        Ax = mpf_add(alpha, a, wp)
681        # case a <= 1
682        if not am[0]:
683            # c = b*b/(r + (a+1)); d = (s + (1-a))
684            # alpha - a = (1/2)*(c + d)
685            # case n=0: re = atan(sqrt((1/2) * Ax * (c + d))/a)
686            # case n=1: re = atan(a/sqrt((1/2) * Ax * (c + d)))
687            c = mpf_div(b2, mpf_add(r, ap, wp), wp)
688            d = mpf_add(s, am, wp)
689            re = mpf_shift(mpf_mul(Ax, mpf_add(c, d, wp), wp), -1)
690            if n == 0:
691                re = mpf_atan(mpf_div(mpf_sqrt(re, wp), a, wp), wp)
692            else:
693                re = mpf_atan(mpf_div(a, mpf_sqrt(re, wp), wp), wp)
694        else:
695            # c = Ax/(r + (a+1)); d = Ax/(s - (1-a))
696            # alpha - a = (1/2)*(c + d)
697            # case n = 0: re = atan(b*sqrt(c + d)/2/a)
698            # case n = 1: re = atan(a/(b*sqrt(c + d)/2)
699            c = mpf_div(Ax, mpf_add(r, ap, wp), wp)
700            d = mpf_div(Ax, mpf_sub(s, am, wp), wp)
701            re = mpf_shift(mpf_add(c, d, wp), -1)
702            re = mpf_mul(b, mpf_sqrt(re, wp), wp)
703            if n == 0:
704                re = mpf_atan(mpf_div(re, a, wp), wp)
705            else:
706                re = mpf_atan(mpf_div(a, re, wp), wp)
707    # to compute alpha + sqrt(alpha**2 - 1), if alpha <= alpha_crossover
708    # replace it with 1 + Am1 + sqrt(Am1*(alpha+1)))
709    # where Am1 = alpha -1
710    # if alpha <= alpha_crossover:
711    if not mpf_sub(alpha_crossover, alpha, wp)[0]:
712        c1 = mpf_div(b2, mpf_add(r, ap, wp), wp)
713        # case a < 1
714        if mpf_neg(am)[0]:
715            # Am1 = (1/2) * (b*b/(r + (a+1)) + b*b/(s + (1-a))
716            c2 = mpf_add(s, am, wp)
717            c2 = mpf_div(b2, c2, wp)
718            Am1 = mpf_shift(mpf_add(c1, c2, wp), -1)
719        else:
720            # Am1 = (1/2) * (b*b/(r + (a+1)) + (s - (1-a)))
721            c2 = mpf_sub(s, am, wp)
722            Am1 = mpf_shift(mpf_add(c1, c2, wp), -1)
723        # im = log(1 + Am1 + sqrt(Am1*(alpha+1)))
724        im = mpf_mul(Am1, mpf_add(alpha, fone, wp), wp)
725        im = mpf_log(mpf_add(fone, mpf_add(Am1, mpf_sqrt(im, wp), wp), wp), wp)
726    else:
727        # im = log(alpha + sqrt(alpha*alpha - 1))
728        im = mpf_sqrt(mpf_sub(mpf_mul(alpha, alpha, wp), fone, wp), wp)
729        im = mpf_log(mpf_add(alpha, im, wp), wp)
730    if asign:
731        if n == 0:
732            re = mpf_sub(mpf_pi(wp), re, wp)
733        else:
734            re = mpf_neg(re)
735    if not bsign and n == 0:
736        im = mpf_neg(im)
737    if bsign and n == 1:
738        im = mpf_neg(im)
739    re = normalize(re[0], re[1], re[2], re[3], prec, rnd)
740    im = normalize(im[0], im[1], im[2], im[3], prec, rnd)
741    return re, im
742
743def mpc_acos(z, prec, rnd=round_fast):
744    return acos_asin(z, prec, rnd, 0)
745
746def mpc_asin(z, prec, rnd=round_fast):
747    return acos_asin(z, prec, rnd, 1)
748
749def mpc_asinh(z, prec, rnd=round_fast):
750    # asinh(z) = I * asin(-I z)
751    a, b = z
752    a, b =  mpc_asin((b, mpf_neg(a)), prec, rnd)
753    return mpf_neg(b), a
754
755def mpc_acosh(z, prec, rnd=round_fast):
756    # acosh(z) = -I * acos(z)   for Im(acos(z)) <= 0
757    #            +I * acos(z)   otherwise
758    a, b = mpc_acos(z, prec, rnd)
759    if b[0] or b == fzero:
760        return mpf_neg(b), a
761    else:
762        return b, mpf_neg(a)
763
764def mpc_atanh(z, prec, rnd=round_fast):
765    # atanh(z) = (log(1+z)-log(1-z))/2
766    wp = prec + 15
767    a = mpc_add(z, mpc_one, wp)
768    b = mpc_sub(mpc_one, z, wp)
769    a = mpc_log(a, wp)
770    b = mpc_log(b, wp)
771    v = mpc_shift(mpc_sub(a, b, wp), -1)
772    # Subtraction at infinity gives correct imaginary part but
773    # wrong real part (should be zero)
774    if v[0] == fnan and mpc_is_inf(z):
775        v = (fzero, v[1])
776    return v
777
778def mpc_fibonacci(z, prec, rnd=round_fast):
779    re, im = z
780    if im == fzero:
781        return (mpf_fibonacci(re, prec, rnd), fzero)
782    size = max(abs(re[2]+re[3]), abs(re[2]+re[3]))
783    wp = prec + size + 20
784    a = mpf_phi(wp)
785    b = mpf_add(mpf_shift(a, 1), fnone, wp)
786    u = mpc_pow((a, fzero), z, wp)
787    v = mpc_cos_pi(z, wp)
788    v = mpc_div(v, u, wp)
789    u = mpc_sub(u, v, wp)
790    u = mpc_div_mpf(u, b, prec, rnd)
791    return u
792
793def mpf_expj(x, prec, rnd='f'):
794    raise ComplexResult
795
796def mpc_expj(z, prec, rnd='f'):
797    re, im = z
798    if im == fzero:
799        return mpf_cos_sin(re, prec, rnd)
800    if re == fzero:
801        return mpf_exp(mpf_neg(im), prec, rnd), fzero
802    ey = mpf_exp(mpf_neg(im), prec+10)
803    c, s = mpf_cos_sin(re, prec+10)
804    re = mpf_mul(ey, c, prec, rnd)
805    im = mpf_mul(ey, s, prec, rnd)
806    return re, im
807
808def mpf_expjpi(x, prec, rnd='f'):
809    raise ComplexResult
810
811def mpc_expjpi(z, prec, rnd='f'):
812    re, im = z
813    if im == fzero:
814        return mpf_cos_sin_pi(re, prec, rnd)
815    sign, man, exp, bc = im
816    wp = prec+10
817    if man:
818        wp += max(0, exp+bc)
819    im = mpf_neg(mpf_mul(mpf_pi(wp), im, wp))
820    if re == fzero:
821        return mpf_exp(im, prec, rnd), fzero
822    ey = mpf_exp(im, prec+10)
823    c, s = mpf_cos_sin_pi(re, prec+10)
824    re = mpf_mul(ey, c, prec, rnd)
825    im = mpf_mul(ey, s, prec, rnd)
826    return re, im
827
828
829if BACKEND == 'sage':
830    try:
831        import sage.libs.mpmath.ext_libmp as _lbmp
832        mpc_exp = _lbmp.mpc_exp
833        mpc_sqrt = _lbmp.mpc_sqrt
834    except (ImportError, AttributeError):
835        print("Warning: Sage imports in libmpc failed")
836