"""
Low-level functions for complex arithmetic.
"""

import sys

from .backend import MPZ, MPZ_ZERO, MPZ_ONE, MPZ_TWO, BACKEND

from .libmpf import (\
    round_floor, round_ceiling, round_down, round_up,
    round_nearest, round_fast, bitcount,
    bctable, normalize, normalize1, reciprocal_rnd, rshift, lshift, giant_steps,
    negative_rnd,
    to_str, to_fixed, from_man_exp, from_float, to_float, from_int, to_int,
    fzero, fone, ftwo, fhalf, finf, fninf, fnan, fnone,
    mpf_abs, mpf_pos, mpf_neg, mpf_add, mpf_sub, mpf_mul,
    mpf_div, mpf_mul_int, mpf_shift, mpf_sqrt, mpf_hypot,
    mpf_rdiv_int, mpf_floor, mpf_ceil, mpf_nint, mpf_frac,
    mpf_sign, mpf_hash,
    ComplexResult
)

from .libelefun import (\
    mpf_pi, mpf_exp, mpf_log, mpf_cos_sin, mpf_cosh_sinh, mpf_tan, mpf_pow_int,
    mpf_log_hypot,
    mpf_cos_sin_pi, mpf_phi,
    mpf_cos, mpf_sin, mpf_cos_pi, mpf_sin_pi,
    mpf_atan, mpf_atan2, mpf_cosh, mpf_sinh, mpf_tanh,
    mpf_asin, mpf_acos, mpf_acosh, mpf_nthroot, mpf_fibonacci
)

# An mpc value is a (real, imag) tuple
mpc_one = fone, fzero
mpc_zero = fzero, fzero
mpc_two = ftwo, fzero
mpc_half = (fhalf, fzero)

_infs = (finf, fninf)
_infs_nan = (finf, fninf, fnan)

def mpc_is_inf(z):
    """Check if either real or imaginary part is infinite"""
    re, im = z
    if re in _infs: return True
    if im in _infs: return True
    return False

def mpc_is_infnan(z):
    """Check if either real or imaginary part is infinite or nan"""
    re, im = z
    if re in _infs_nan: return True
    if im in _infs_nan: return True
    return False

def mpc_to_str(z, dps, **kwargs):
    re, im = z
    rs = to_str(re, dps)
    if im[0]:
        return rs + " - " + to_str(mpf_neg(im), dps, **kwargs) + "j"
    else:
        return rs + " + " + to_str(im, dps, **kwargs) + "j"

def mpc_to_complex(z, strict=False, rnd=round_fast):
    re, im = z
    return complex(to_float(re, strict, rnd), to_float(im, strict, rnd))

def mpc_hash(z):
    if sys.version >= "3.2":
        re, im = z
        h = mpf_hash(re) + sys.hash_info.imag * mpf_hash(im)
        # Need to reduce either module 2^32 or 2^64
        h = h % (2**sys.hash_info.width)
        return int(h)
    else:
        try:
            return hash(mpc_to_complex(z, strict=True))
        except OverflowError:
            return hash(z)

def mpc_conjugate(z, prec, rnd=round_fast):
    re, im = z
    return re, mpf_neg(im, prec, rnd)

def mpc_is_nonzero(z):
    return z != mpc_zero

def mpc_add(z, w, prec, rnd=round_fast):
    a, b = z
    c, d = w
    return mpf_add(a, c, prec, rnd), mpf_add(b, d, prec, rnd)

def mpc_add_mpf(z, x, prec, rnd=round_fast):
    a, b = z
    return mpf_add(a, x, prec, rnd), b

def mpc_sub(z, w, prec=0, rnd=round_fast):
    a, b = z
    c, d = w
    return mpf_sub(a, c, prec, rnd), mpf_sub(b, d, prec, rnd)

def mpc_sub_mpf(z, p, prec=0, rnd=round_fast):
    a, b = z
    return mpf_sub(a, p, prec, rnd), b

def mpc_pos(z, prec, rnd=round_fast):
    a, b = z
    return mpf_pos(a, prec, rnd), mpf_pos(b, prec, rnd)

def mpc_neg(z, prec=None, rnd=round_fast):
    a, b = z
    return mpf_neg(a, prec, rnd), mpf_neg(b, prec, rnd)

def mpc_shift(z, n):
    a, b = z
    return mpf_shift(a, n), mpf_shift(b, n)

def mpc_abs(z, prec, rnd=round_fast):
    """Absolute value of a complex number, |a+bi|.
    Returns an mpf value."""
    a, b = z
    return mpf_hypot(a, b, prec, rnd)

def mpc_arg(z, prec, rnd=round_fast):
    """Argument of a complex number. Returns an mpf value."""
    a, b = z
    return mpf_atan2(b, a, prec, rnd)

def mpc_floor(z, prec, rnd=round_fast):
    a, b = z
    return mpf_floor(a, prec, rnd), mpf_floor(b, prec, rnd)

def mpc_ceil(z, prec, rnd=round_fast):
    a, b = z
    return mpf_ceil(a, prec, rnd), mpf_ceil(b, prec, rnd)

def mpc_nint(z, prec, rnd=round_fast):
    a, b = z
    return mpf_nint(a, prec, rnd), mpf_nint(b, prec, rnd)

def mpc_frac(z, prec, rnd=round_fast):
    a, b = z
    return mpf_frac(a, prec, rnd), mpf_frac(b, prec, rnd)

def mpc_mul(z, w, prec, rnd=round_fast):
    """
    Complex multiplication.

    Returns the real and imaginary part of (a+bi)*(c+di), rounded to
    the specified precision. The rounding mode applies to the real and
    imaginary parts separately.
    """
    a, b = z
    c, d = w
    p = mpf_mul(a, c)
    q = mpf_mul(b, d)
    r = mpf_mul(a, d)
    s = mpf_mul(b, c)
    re = mpf_sub(p, q, prec, rnd)
    im = mpf_add(r, s, prec, rnd)
    return re, im

def mpc_square(z, prec, rnd=round_fast):
    # (a+b*I)**2 == a**2 - b**2 + 2*I*a*b
    a, b = z
    p = mpf_mul(a,a)
    q = mpf_mul(b,b)
    r = mpf_mul(a,b, prec, rnd)
    re = mpf_sub(p, q, prec, rnd)
    im = mpf_shift(r, 1)
    return re, im

def mpc_mul_mpf(z, p, prec, rnd=round_fast):
    a, b = z
    re = mpf_mul(a, p, prec, rnd)
    im = mpf_mul(b, p, prec, rnd)
    return re, im

def mpc_mul_imag_mpf(z, x, prec, rnd=round_fast):
    """
    Multiply the mpc value z by I*x where x is an mpf value.
    """
    a, b = z
    re = mpf_neg(mpf_mul(b, x, prec, rnd))
    im = mpf_mul(a, x, prec, rnd)
    return re, im

def mpc_mul_int(z, n, prec, rnd=round_fast):
    a, b = z
    re = mpf_mul_int(a, n, prec, rnd)
    im = mpf_mul_int(b, n, prec, rnd)
    return re, im

def mpc_div(z, w, prec, rnd=round_fast):
    a, b = z
    c, d = w
    wp = prec + 10
    # mag = c*c + d*d
    mag = mpf_add(mpf_mul(c, c), mpf_mul(d, d), wp)
    # (a*c+b*d)/mag, (b*c-a*d)/mag
    t = mpf_add(mpf_mul(a,c), mpf_mul(b,d), wp)
    u = mpf_sub(mpf_mul(b,c), mpf_mul(a,d), wp)
    return mpf_div(t,mag,prec,rnd), mpf_div(u,mag,prec,rnd)

def mpc_div_mpf(z, p, prec, rnd=round_fast):
    """Calculate z/p where p is real"""
    a, b = z
    re = mpf_div(a, p, prec, rnd)
    im = mpf_div(b, p, prec, rnd)
    return re, im

def mpc_reciprocal(z, prec, rnd=round_fast):
    """Calculate 1/z efficiently"""
    a, b = z
    m = mpf_add(mpf_mul(a,a),mpf_mul(b,b),prec+10)
    re = mpf_div(a, m, prec, rnd)
    im = mpf_neg(mpf_div(b, m, prec, rnd))
    return re, im

def mpc_mpf_div(p, z, prec, rnd=round_fast):
    """Calculate p/z where p is real efficiently"""
    a, b = z
    m = mpf_add(mpf_mul(a,a),mpf_mul(b,b), prec+10)
    re = mpf_div(mpf_mul(a,p), m, prec, rnd)
    im = mpf_div(mpf_neg(mpf_mul(b,p)), m, prec, rnd)
    return re, im

def complex_int_pow(a, b, n):
    """Complex integer power: computes (a+b*I)**n exactly for
    nonnegative n (a and b must be Python ints)."""
    wre = 1
    wim = 0
    while n:
        if n & 1:
            wre, wim = wre*a - wim*b, wim*a + wre*b
            n -= 1
        a, b = a*a - b*b, 2*a*b
        n //= 2
    return wre, wim

def mpc_pow(z, w, prec, rnd=round_fast):
    if w[1] == fzero:
        return mpc_pow_mpf(z, w[0], prec, rnd)
    return mpc_exp(mpc_mul(mpc_log(z, prec+10), w, prec+10), prec, rnd)

def mpc_pow_mpf(z, p, prec, rnd=round_fast):
    psign, pman, pexp, pbc = p
    if pexp >= 0:
        return mpc_pow_int(z, (-1)**psign * (pman<<pexp), prec, rnd)
    if pexp == -1:
        sqrtz = mpc_sqrt(z, prec+10)
        return mpc_pow_int(sqrtz, (-1)**psign * pman, prec, rnd)
    return mpc_exp(mpc_mul_mpf(mpc_log(z, prec+10), p, prec+10), prec, rnd)

def mpc_pow_int(z, n, prec, rnd=round_fast):
    a, b = z
    if b == fzero:
        return mpf_pow_int(a, n, prec, rnd), fzero
    if a == fzero:
        v = mpf_pow_int(b, n, prec, rnd)
        n %= 4
        if n == 0:
            return v, fzero
        elif n == 1:
            return fzero, v
        elif n == 2:
            return mpf_neg(v), fzero
        elif n == 3:
            return fzero, mpf_neg(v)
    if n == 0: return mpc_one
    if n == 1: return mpc_pos(z, prec, rnd)
    if n == 2: return mpc_square(z, prec, rnd)
    if n == -1: return mpc_reciprocal(z, prec, rnd)
    if n < 0: return mpc_reciprocal(mpc_pow_int(z, -n, prec+4), prec, rnd)
    asign, aman, aexp, abc = a
    bsign, bman, bexp, bbc = b
    if asign: aman = -aman
    if bsign: bman = -bman
    de = aexp - bexp
    abs_de = abs(de)
    exact_size = n*(abs_de + max(abc, bbc))
    if exact_size < 10000:
        if de > 0:
            aman <<= de
            aexp = bexp
        else:
            bman <<= (-de)
            bexp = aexp
        re, im = complex_int_pow(aman, bman, n)
        re = from_man_exp(re, int(n*aexp), prec, rnd)
        im = from_man_exp(im, int(n*bexp), prec, rnd)
        return re, im
    return mpc_exp(mpc_mul_int(mpc_log(z, prec+10), n, prec+10), prec, rnd)

def mpc_sqrt(z, prec, rnd=round_fast):
    """Complex square root (principal branch).

    We have sqrt(a+bi) = sqrt((r+a)/2) + b/sqrt(2*(r+a))*i where
    r = abs(a+bi), when a+bi is not a negative real number."""
    a, b = z
    if b == fzero:
        if a == fzero:
            return (a, b)
        # When a+bi is a negative real number, we get a real sqrt times i
        if a[0]:
            im = mpf_sqrt(mpf_neg(a), prec, rnd)
            return (fzero, im)
        else:
            re = mpf_sqrt(a, prec, rnd)
            return (re, fzero)
    wp = prec+20
    if not a[0]:                               # case a positive
        t  = mpf_add(mpc_abs((a, b), wp), a, wp)  # t = abs(a+bi) + a
        u = mpf_shift(t, -1)                      # u = t/2
        re = mpf_sqrt(u, prec, rnd)               # re = sqrt(u)
        v = mpf_shift(t, 1)                       # v = 2*t
        w  = mpf_sqrt(v, wp)                      # w = sqrt(v)
        im = mpf_div(b, w, prec, rnd)             # im = b / w
    else:                                      # case a negative
        t = mpf_sub(mpc_abs((a, b), wp), a, wp)   # t = abs(a+bi) - a
        u = mpf_shift(t, -1)                      # u = t/2
        im = mpf_sqrt(u, prec, rnd)               # im = sqrt(u)
        v = mpf_shift(t, 1)                       # v = 2*t
        w  = mpf_sqrt(v, wp)                      # w = sqrt(v)
        re = mpf_div(b, w, prec, rnd)             # re = b/w
        if b[0]:
            re = mpf_neg(re)
            im = mpf_neg(im)
    return re, im

def mpc_nthroot_fixed(a, b, n, prec):
    # a, b signed integers at fixed precision prec
    start = 50
    a1 = int(rshift(a, prec - n*start))
    b1 = int(rshift(b, prec - n*start))
    try:
        r = (a1 + 1j * b1)**(1.0/n)
        re = r.real
        im = r.imag
        re = MPZ(int(re))
        im = MPZ(int(im))
    except OverflowError:
        a1 = from_int(a1, start)
        b1 = from_int(b1, start)
        fn = from_int(n)
        nth = mpf_rdiv_int(1, fn, start)
        re, im = mpc_pow((a1, b1), (nth, fzero), start)
        re = to_int(re)
        im = to_int(im)
    extra = 10
    prevp = start
    extra1 = n
    for p in giant_steps(start, prec+extra):
        # this is slow for large n, unlike int_pow_fixed
        re2, im2 = complex_int_pow(re, im, n-1)
        re2 = rshift(re2, (n-1)*prevp - p - extra1)
        im2 = rshift(im2, (n-1)*prevp - p - extra1)
        r4 = (re2*re2 + im2*im2) >> (p + extra1)
        ap = rshift(a, prec - p)
        bp = rshift(b, prec - p)
        rec = (ap * re2 + bp * im2) >> p
        imc = (-ap * im2 + bp * re2) >> p
        reb = (rec << p) // r4
        imb = (imc << p) // r4
        re = (reb + (n-1)*lshift(re, p-prevp))//n
        im = (imb + (n-1)*lshift(im, p-prevp))//n
        prevp = p
    return re, im

def mpc_nthroot(z, n, prec, rnd=round_fast):
    """
    Complex n-th root.

    Use Newton method as in the real case when it is faster,
    otherwise use z**(1/n)
    """
    a, b = z
    if a[0] == 0 and b == fzero:
        re = mpf_nthroot(a, n, prec, rnd)
        return (re, fzero)
    if n < 2:
        if n == 0:
            return mpc_one
        if n == 1:
            return mpc_pos((a, b), prec, rnd)
        if n == -1:
            return mpc_div(mpc_one, (a, b), prec, rnd)
        inverse = mpc_nthroot((a, b), -n, prec+5, reciprocal_rnd[rnd])
        return mpc_div(mpc_one, inverse, prec, rnd)
    if n <= 20:
        prec2 = int(1.2 * (prec + 10))
        asign, aman, aexp, abc = a
        bsign, bman, bexp, bbc = b
        pf = mpc_abs((a,b), prec)
        if pf[-2] + pf[-1] > -10  and pf[-2] + pf[-1] < prec:
            af = to_fixed(a, prec2)
            bf = to_fixed(b, prec2)
            re, im = mpc_nthroot_fixed(af, bf, n, prec2)
            extra = 10
            re = from_man_exp(re, -prec2-extra, prec2, rnd)
            im = from_man_exp(im, -prec2-extra, prec2, rnd)
            return re, im
    fn = from_int(n)
    prec2 = prec+10 + 10
    nth = mpf_rdiv_int(1, fn, prec2)
    re, im = mpc_pow((a, b), (nth, fzero), prec2, rnd)
    re = normalize(re[0], re[1], re[2], re[3], prec, rnd)
    im = normalize(im[0], im[1], im[2], im[3], prec, rnd)
    return re, im

def mpc_cbrt(z, prec, rnd=round_fast):
    """
    Complex cubic root.
    """
    return mpc_nthroot(z, 3, prec, rnd)

def mpc_exp(z, prec, rnd=round_fast):
    """
    Complex exponential function.

    We use the direct formula exp(a+bi) = exp(a) * (cos(b) + sin(b)*i)
    for the computation. This formula is very nice because it is
    pefectly stable; since we just do real multiplications, the only
    numerical errors that can creep in are single-ulp rounding errors.

    The formula is efficient since mpmath's real exp is quite fast and
    since we can compute cos and sin simultaneously.

    It is no problem if a and b are large; if the implementations of
    exp/cos/sin are accurate and efficient for all real numbers, then
    so is this function for all complex numbers.
    """
    a, b = z
    if a == fzero:
        return mpf_cos_sin(b, prec, rnd)
    if b == fzero:
        return mpf_exp(a, prec, rnd), fzero
    mag = mpf_exp(a, prec+4, rnd)
    c, s = mpf_cos_sin(b, prec+4, rnd)
    re = mpf_mul(mag, c, prec, rnd)
    im = mpf_mul(mag, s, prec, rnd)
    return re, im

def mpc_log(z, prec, rnd=round_fast):
    re = mpf_log_hypot(z[0], z[1], prec, rnd)
    im = mpc_arg(z, prec, rnd)
    return re, im

def mpc_cos(z, prec, rnd=round_fast):
    """Complex cosine. The formula used is cos(a+bi) = cos(a)*cosh(b) -
    sin(a)*sinh(b)*i.

    The same comments apply as for the complex exp: only real
    multiplications are pewrormed, so no cancellation errors are
    possible. The formula is also efficient since we can compute both
    pairs (cos, sin) and (cosh, sinh) in single stwps."""
    a, b = z
    if b == fzero:
        return mpf_cos(a, prec, rnd), fzero
    if a == fzero:
        return mpf_cosh(b, prec, rnd), fzero
    wp = prec + 6
    c, s = mpf_cos_sin(a, wp)
    ch, sh = mpf_cosh_sinh(b, wp)
    re = mpf_mul(c, ch, prec, rnd)
    im = mpf_mul(s, sh, prec, rnd)
    return re, mpf_neg(im)

def mpc_sin(z, prec, rnd=round_fast):
    """Complex sine. We have sin(a+bi) = sin(a)*cosh(b) +
    cos(a)*sinh(b)*i. See the docstring for mpc_cos for additional
    comments."""
    a, b = z
    if b == fzero:
        return mpf_sin(a, prec, rnd), fzero
    if a == fzero:
        return fzero, mpf_sinh(b, prec, rnd)
    wp = prec + 6
    c, s = mpf_cos_sin(a, wp)
    ch, sh = mpf_cosh_sinh(b, wp)
    re = mpf_mul(s, ch, prec, rnd)
    im = mpf_mul(c, sh, prec, rnd)
    return re, im

def mpc_tan(z, prec, rnd=round_fast):
    """Complex tangent. Computed as tan(a+bi) = sin(2a)/M + sinh(2b)/M*i
    where M = cos(2a) + cosh(2b)."""
    a, b = z
    asign, aman, aexp, abc = a
    bsign, bman, bexp, bbc = b
    if b == fzero: return mpf_tan(a, prec, rnd), fzero
    if a == fzero: return fzero, mpf_tanh(b, prec, rnd)
    wp = prec + 15
    a = mpf_shift(a, 1)
    b = mpf_shift(b, 1)
    c, s = mpf_cos_sin(a, wp)
    ch, sh = mpf_cosh_sinh(b, wp)
    # TODO: handle cancellation when c ~=  -1 and ch ~= 1
    mag = mpf_add(c, ch, wp)
    re = mpf_div(s, mag, prec, rnd)
    im = mpf_div(sh, mag, prec, rnd)
    return re, im

def mpc_cos_pi(z, prec, rnd=round_fast):
    a, b = z
    if b == fzero:
        return mpf_cos_pi(a, prec, rnd), fzero
    b = mpf_mul(b, mpf_pi(prec+5), prec+5)
    if a == fzero:
        return mpf_cosh(b, prec, rnd), fzero
    wp = prec + 6
    c, s = mpf_cos_sin_pi(a, wp)
    ch, sh = mpf_cosh_sinh(b, wp)
    re = mpf_mul(c, ch, prec, rnd)
    im = mpf_mul(s, sh, prec, rnd)
    return re, mpf_neg(im)

def mpc_sin_pi(z, prec, rnd=round_fast):
    a, b = z
    if b == fzero:
        return mpf_sin_pi(a, prec, rnd), fzero
    b = mpf_mul(b, mpf_pi(prec+5), prec+5)
    if a == fzero:
        return fzero, mpf_sinh(b, prec, rnd)
    wp = prec + 6
    c, s = mpf_cos_sin_pi(a, wp)
    ch, sh = mpf_cosh_sinh(b, wp)
    re = mpf_mul(s, ch, prec, rnd)
    im = mpf_mul(c, sh, prec, rnd)
    return re, im

def mpc_cos_sin(z, prec, rnd=round_fast):
    a, b = z
    if a == fzero:
        ch, sh = mpf_cosh_sinh(b, prec, rnd)
        return (ch, fzero), (fzero, sh)
    if b == fzero:
        c, s = mpf_cos_sin(a, prec, rnd)
        return (c, fzero), (s, fzero)
    wp = prec + 6
    c, s = mpf_cos_sin(a, wp)
    ch, sh = mpf_cosh_sinh(b, wp)
    cre = mpf_mul(c, ch, prec, rnd)
    cim = mpf_mul(s, sh, prec, rnd)
    sre = mpf_mul(s, ch, prec, rnd)
    sim = mpf_mul(c, sh, prec, rnd)
    return (cre, mpf_neg(cim)), (sre, sim)

def mpc_cos_sin_pi(z, prec, rnd=round_fast):
    a, b = z
    if b == fzero:
        c, s = mpf_cos_sin_pi(a, prec, rnd)
        return (c, fzero), (s, fzero)
    b = mpf_mul(b, mpf_pi(prec+5), prec+5)
    if a == fzero:
        ch, sh = mpf_cosh_sinh(b, prec, rnd)
        return (ch, fzero), (fzero, sh)
    wp = prec + 6
    c, s = mpf_cos_sin_pi(a, wp)
    ch, sh = mpf_cosh_sinh(b, wp)
    cre = mpf_mul(c, ch, prec, rnd)
    cim = mpf_mul(s, sh, prec, rnd)
    sre = mpf_mul(s, ch, prec, rnd)
    sim = mpf_mul(c, sh, prec, rnd)
    return (cre, mpf_neg(cim)), (sre, sim)

def mpc_cosh(z, prec, rnd=round_fast):
    """Complex hyperbolic cosine. Computed as cosh(z) = cos(z*i)."""
    a, b = z
    return mpc_cos((b, mpf_neg(a)), prec, rnd)

def mpc_sinh(z, prec, rnd=round_fast):
    """Complex hyperbolic sine. Computed as sinh(z) = -i*sin(z*i)."""
    a, b = z
    b, a = mpc_sin((b, a), prec, rnd)
    return a, b

def mpc_tanh(z, prec, rnd=round_fast):
    """Complex hyperbolic tangent. Computed as tanh(z) = -i*tan(z*i)."""
    a, b = z
    b, a = mpc_tan((b, a), prec, rnd)
    return a, b

# TODO: avoid loss of accuracy
def mpc_atan(z, prec, rnd=round_fast):
    a, b = z
    # atan(z) = (I/2)*(log(1-I*z) - log(1+I*z))
    # x = 1-I*z = 1 + b - I*a
    # y = 1+I*z = 1 - b + I*a
    wp = prec + 15
    x = mpf_add(fone, b, wp), mpf_neg(a)
    y = mpf_sub(fone, b, wp), a
    l1 = mpc_log(x, wp)
    l2 = mpc_log(y, wp)
    a, b = mpc_sub(l1, l2, prec, rnd)
    # (I/2) * (a+b*I) = (-b/2 + a/2*I)
    v = mpf_neg(mpf_shift(b,-1)), mpf_shift(a,-1)
    # Subtraction at infinity gives correct real part but
    # wrong imaginary part (should be zero)
    if v[1] == fnan and mpc_is_inf(z):
        v = (v[0], fzero)
    return v

beta_crossover = from_float(0.6417)
alpha_crossover = from_float(1.5)

def acos_asin(z, prec, rnd, n):
    """ complex acos for n = 0, asin for n = 1
    The algorithm is described in
    T.E. Hull, T.F. Fairgrieve and P.T.P. Tang
    'Implementing the Complex Arcsine and Arcosine Functions
    using Exception Handling',
    ACM Trans. on Math. Software Vol. 23 (1997), p299
    The complex acos and asin can be defined as
    acos(z) = acos(beta) - I*sign(a)* log(alpha + sqrt(alpha**2 -1))
    asin(z) = asin(beta) + I*sign(a)* log(alpha + sqrt(alpha**2 -1))
    where z = a + I*b
    alpha = (1/2)*(r + s); beta = (1/2)*(r - s) = a/alpha
    r = sqrt((a+1)**2 + y**2); s = sqrt((a-1)**2 + y**2)
    These expressions are rewritten in different ways in different
    regions, delimited by two crossovers alpha_crossover and beta_crossover,
    and by abs(a) <= 1, in order to improve the numerical accuracy.
    """
    a, b = z
    wp = prec + 10
    # special cases with real argument
    if b == fzero:
        am = mpf_sub(fone, mpf_abs(a), wp)
        # case abs(a) <= 1
        if not am[0]:
            if n == 0:
                return mpf_acos(a, prec, rnd), fzero
            else:
                return mpf_asin(a, prec, rnd), fzero
        # cases abs(a) > 1
        else:
            # case a < -1
            if a[0]:
                pi = mpf_pi(prec, rnd)
                c = mpf_acosh(mpf_neg(a), prec, rnd)
                if n == 0:
                    return pi, mpf_neg(c)
                else:
                    return mpf_neg(mpf_shift(pi, -1)), c
            # case a > 1
            else:
                c = mpf_acosh(a, prec, rnd)
                if n == 0:
                    return fzero, c
                else:
                    pi = mpf_pi(prec, rnd)
                    return mpf_shift(pi, -1), mpf_neg(c)
    asign = bsign = 0
    if a[0]:
        a = mpf_neg(a)
        asign = 1
    if b[0]:
        b = mpf_neg(b)
        bsign = 1
    am = mpf_sub(fone, a, wp)
    ap = mpf_add(fone, a, wp)
    r = mpf_hypot(ap, b, wp)
    s = mpf_hypot(am, b, wp)
    alpha = mpf_shift(mpf_add(r, s, wp), -1)
    beta = mpf_div(a, alpha, wp)
    b2 = mpf_mul(b,b, wp)
    # case beta <= beta_crossover
    if not mpf_sub(beta_crossover, beta, wp)[0]:
        if n == 0:
            re = mpf_acos(beta, wp)
        else:
            re = mpf_asin(beta, wp)
    else:
        # to compute the real part in this region use the identity
        # asin(beta) = atan(beta/sqrt(1-beta**2))
        # beta/sqrt(1-beta**2) = (alpha + a) * (alpha - a)
        # alpha + a is numerically accurate; alpha - a can have
        # cancellations leading to numerical inaccuracies, so rewrite
        # it in differente ways according to the region
        Ax = mpf_add(alpha, a, wp)
        # case a <= 1
        if not am[0]:
            # c = b*b/(r + (a+1)); d = (s + (1-a))
            # alpha - a = (1/2)*(c + d)
            # case n=0: re = atan(sqrt((1/2) * Ax * (c + d))/a)
            # case n=1: re = atan(a/sqrt((1/2) * Ax * (c + d)))
            c = mpf_div(b2, mpf_add(r, ap, wp), wp)
            d = mpf_add(s, am, wp)
            re = mpf_shift(mpf_mul(Ax, mpf_add(c, d, wp), wp), -1)
            if n == 0:
                re = mpf_atan(mpf_div(mpf_sqrt(re, wp), a, wp), wp)
            else:
                re = mpf_atan(mpf_div(a, mpf_sqrt(re, wp), wp), wp)
        else:
            # c = Ax/(r + (a+1)); d = Ax/(s - (1-a))
            # alpha - a = (1/2)*(c + d)
            # case n = 0: re = atan(b*sqrt(c + d)/2/a)
            # case n = 1: re = atan(a/(b*sqrt(c + d)/2)
            c = mpf_div(Ax, mpf_add(r, ap, wp), wp)
            d = mpf_div(Ax, mpf_sub(s, am, wp), wp)
            re = mpf_shift(mpf_add(c, d, wp), -1)
            re = mpf_mul(b, mpf_sqrt(re, wp), wp)
            if n == 0:
                re = mpf_atan(mpf_div(re, a, wp), wp)
            else:
                re = mpf_atan(mpf_div(a, re, wp), wp)
    # to compute alpha + sqrt(alpha**2 - 1), if alpha <= alpha_crossover
    # replace it with 1 + Am1 + sqrt(Am1*(alpha+1)))
    # where Am1 = alpha -1
    # if alpha <= alpha_crossover:
    if not mpf_sub(alpha_crossover, alpha, wp)[0]:
        c1 = mpf_div(b2, mpf_add(r, ap, wp), wp)
        # case a < 1
        if mpf_neg(am)[0]:
            # Am1 = (1/2) * (b*b/(r + (a+1)) + b*b/(s + (1-a))
            c2 = mpf_add(s, am, wp)
            c2 = mpf_div(b2, c2, wp)
            Am1 = mpf_shift(mpf_add(c1, c2, wp), -1)
        else:
            # Am1 = (1/2) * (b*b/(r + (a+1)) + (s - (1-a)))
            c2 = mpf_sub(s, am, wp)
            Am1 = mpf_shift(mpf_add(c1, c2, wp), -1)
        # im = log(1 + Am1 + sqrt(Am1*(alpha+1)))
        im = mpf_mul(Am1, mpf_add(alpha, fone, wp), wp)
        im = mpf_log(mpf_add(fone, mpf_add(Am1, mpf_sqrt(im, wp), wp), wp), wp)
    else:
        # im = log(alpha + sqrt(alpha*alpha - 1))
        im = mpf_sqrt(mpf_sub(mpf_mul(alpha, alpha, wp), fone, wp), wp)
        im = mpf_log(mpf_add(alpha, im, wp), wp)
    if asign:
        if n == 0:
            re = mpf_sub(mpf_pi(wp), re, wp)
        else:
            re = mpf_neg(re)
    if not bsign and n == 0:
        im = mpf_neg(im)
    if bsign and n == 1:
        im = mpf_neg(im)
    re = normalize(re[0], re[1], re[2], re[3], prec, rnd)
    im = normalize(im[0], im[1], im[2], im[3], prec, rnd)
    return re, im

def mpc_acos(z, prec, rnd=round_fast):
    return acos_asin(z, prec, rnd, 0)

def mpc_asin(z, prec, rnd=round_fast):
    return acos_asin(z, prec, rnd, 1)

def mpc_asinh(z, prec, rnd=round_fast):
    # asinh(z) = I * asin(-I z)
    a, b = z
    a, b =  mpc_asin((b, mpf_neg(a)), prec, rnd)
    return mpf_neg(b), a

def mpc_acosh(z, prec, rnd=round_fast):
    # acosh(z) = -I * acos(z)   for Im(acos(z)) <= 0
    #            +I * acos(z)   otherwise
    a, b = mpc_acos(z, prec, rnd)
    if b[0] or b == fzero:
        return mpf_neg(b), a
    else:
        return b, mpf_neg(a)

def mpc_atanh(z, prec, rnd=round_fast):
    # atanh(z) = (log(1+z)-log(1-z))/2
    wp = prec + 15
    a = mpc_add(z, mpc_one, wp)
    b = mpc_sub(mpc_one, z, wp)
    a = mpc_log(a, wp)
    b = mpc_log(b, wp)
    v = mpc_shift(mpc_sub(a, b, wp), -1)
    # Subtraction at infinity gives correct imaginary part but
    # wrong real part (should be zero)
    if v[0] == fnan and mpc_is_inf(z):
        v = (fzero, v[1])
    return v

def mpc_fibonacci(z, prec, rnd=round_fast):
    re, im = z
    if im == fzero:
        return (mpf_fibonacci(re, prec, rnd), fzero)
    size = max(abs(re[2]+re[3]), abs(re[2]+re[3]))
    wp = prec + size + 20
    a = mpf_phi(wp)
    b = mpf_add(mpf_shift(a, 1), fnone, wp)
    u = mpc_pow((a, fzero), z, wp)
    v = mpc_cos_pi(z, wp)
    v = mpc_div(v, u, wp)
    u = mpc_sub(u, v, wp)
    u = mpc_div_mpf(u, b, prec, rnd)
    return u

def mpf_expj(x, prec, rnd='f'):
    raise ComplexResult

def mpc_expj(z, prec, rnd='f'):
    re, im = z
    if im == fzero:
        return mpf_cos_sin(re, prec, rnd)
    if re == fzero:
        return mpf_exp(mpf_neg(im), prec, rnd), fzero
    ey = mpf_exp(mpf_neg(im), prec+10)
    c, s = mpf_cos_sin(re, prec+10)
    re = mpf_mul(ey, c, prec, rnd)
    im = mpf_mul(ey, s, prec, rnd)
    return re, im

def mpf_expjpi(x, prec, rnd='f'):
    raise ComplexResult

def mpc_expjpi(z, prec, rnd='f'):
    re, im = z
    if im == fzero:
        return mpf_cos_sin_pi(re, prec, rnd)
    sign, man, exp, bc = im
    wp = prec+10
    if man:
        wp += max(0, exp+bc)
    im = mpf_neg(mpf_mul(mpf_pi(wp), im, wp))
    if re == fzero:
        return mpf_exp(im, prec, rnd), fzero
    ey = mpf_exp(im, prec+10)
    c, s = mpf_cos_sin_pi(re, prec+10)
    re = mpf_mul(ey, c, prec, rnd)
    im = mpf_mul(ey, s, prec, rnd)
    return re, im

if BACKEND == 'sage':
    try:
        import sage.libs.mpmath.ext_libmp as _lbmp
        mpc_exp = _lbmp.mpc_exp
        mpc_sqrt = _lbmp.mpc_sqrt
    except (ImportError, AttributeError):
        print("Warning: Sage imports in libmpc failed")


# coding: utf-8
# Copyright (c) Materials Virtual Lab
# Distributed under the terms of the BSD License.

import re
import os
import glob
import itertools
import subprocess
from collections import OrderedDict

import numpy as np
import pandas as pd
from monty.io import zopen
from monty.os.path import which
from monty.tempfile import ScratchDir
from monty.serialization import loadfn
from pymatgen import Structure, Lattice, Element
from pymatgen.core import units

from mlearn.potentials import Potential
from mlearn.data import pool_from, convert_docs
from mlearn.potentials.lammps.calcs import EnergyForceStress

module_dir = os.path.dirname(__file__)
NNinput_params = loadfn(os.path.join(module_dir, 'params', 'NNinput.json'))

class NNPotential(Potential):
    """
    This class implements Neural Network Potential.
    """
    bohr_to_angstrom = units.bohr_to_angstrom
    eV_to_Ha = units.eV_to_Ha
    pair_style = 'pair_style        nnp dir "./" showew no showewsum 0 ' \
                 'maxew 10000000 resetew yes cflength 1.8897261328 cfenergy 0.0367493254'
    pair_coeff = 'pair_coeff        * * {}'
    def __init__(self, name=None, param=None, weight_param=None, scaling_param=None):
        """

        Args:
            name (str): Name of force field.
        """
        self.name = name if name else "NNPotential"
        self.specie = None
        self.weights = []
        self.bs = []
        self.atom_energy = None
        self.normalized_nodes = None
        self.epochs = None
        self.param = param if param else {}
        self.weight_param = weight_param if weight_param else None
        self.scaling_param = scaling_param if scaling_param else None
        self.fitted = False

    def _line_up(self, structure, energy, forces, virial_stress):
        """
        Convert input structure, energy, forces, virial_stress to
        proper configuration format for RuNNer usage. Note that
        RuNNer takes bohr as length unit and Hatree as energy unit.

        Args:
            structure (Structure): Pymatgen Structure object.
            energy (float): DFT-calculated energy of the system.
            forces (list): The forces should have dimension
                (num_atoms, 3).
            virial_stress (list): stress should has 6 distinct
                elements arranged in order [xx, yy, zz, xy, yz, xz].

        Returns:
        """
        if len(structure.symbol_set) > 1:
            raise ValueError("Structure is not unary.")

        inputs = OrderedDict(Size=structure.num_sites, \
                             SuperCell=structure.lattice, \
                             AtomData=(structure, forces), \
                             Energy=energy, \
                             Stress=virial_stress)

        lines = ['begin']

        if 'SuperCell' in inputs:
            bohr_matrix = inputs['SuperCell'].matrix / self.bohr_to_angstrom
            for vec in bohr_matrix:
                lines.append('lattice {:>15.6f}{:>15.6f}{:>15.6f}'.format(*vec))
        if 'AtomData' in inputs:
            format_float = \
                'atom{:>16.9f}{:>16.9f}{:>16.9f}{:>4s}{:>15.9f}{:>15.9f}{:>15.9f}{:>15.9f}{:>15.9f}'
            for i, (site, force) in enumerate(zip(structure, forces)):
                lines.append(format_float.format(*site.coords / self.bohr_to_angstrom, \
                                site.species_string, 0.0, 0.0,
                                *np.array(force) * self.eV_to_Ha * self.bohr_to_angstrom))
        if 'Energy' in inputs:
            lines.append('energy  {:f}'.format(energy * self.eV_to_Ha))

        lines.append('charge  {:f}'.format(structure.charge))
        lines.append('end')

        return '\n'.join(lines)

    def write_cfgs(self, filename, cfg_pool):

        lines = []
        for dataset in cfg_pool:
            if isinstance(dataset['structure'], dict):
                structure = Structure.from_dict(dataset['structure'])
            else:
                structure = dataset['structure']
            energy = dataset['outputs']['energy']
            forces = dataset['outputs']['forces']
            virial_stress = dataset['outputs']['virial_stress']

            lines.append(self._line_up(structure, energy, forces, virial_stress))

            # dist = np.unique(structure.distance_matrix.ravel())[1]
            # if self.shortest_distance > dist:
            #     self.shortest_distance = dist

        self.specie = Element(structure.symbol_set[0])

        with open(filename, 'w') as f:
            f.write('\n'.join(lines))

        return filename

    def write_input(self, **kwargs):
        """
        Write input.nn file to train the Neural Network Potential.

        Args:
            atom_energy (float): Atomic reference energy.

            kwargs:
                General nnp settings:
                    atom_energy (None): Free atom reference energy.
                    cutoff_type (int): Type of cutoff function. Default to 1
                        (i.e., cosine function).
                    scale_features (int): Determine the method to scale the
                        symmetry function.
                        0: no scaling.
                        1: scale_symmetry_functions.
                        2: center_symmetry_functions.
                        3. scale_symmetry_functions_sigma.
                    scale_min_short (float): Minimum value for scaling.
                        Default to 0.0.
                    scale_max_short (float): Maximum value for scaling.
                        Default to 1.
                    hidden_layers (list): List of the numbers of
                        nodes in each hidden layer.
                    activations (str): Activation function for each hidden layer.
                        't': tanh, 's': logistic, 'p': softplus.
                    normalize_nodes (boolean): Whether to normalize input of nodes.

                Additional settings for training:
                    epoch (int): Number of training epochs.
                    updater_type (int): Weight update method
                        0: gradient Descent, 1: Kalman filter.
                    parallel_mode (int): Training parallelization used.
                        Default to serial mode.
                    update_strategy (int): Update strategy.
                        0: combined, 1: per-element.
                    selection_mode (int): Update candidate selection mode.
                        0: random, 1: sort, 2: threshold
                    test_fraction (float): Fraction of structures kept for
                        testing.
                    force_weight (float): Weight of force updates relative
                        to energy updates. Default to 10.0
                    short_energy_fraction (float): Fraction of energy updates
                        per epoch. Default to 1.0.
                    short_force_fraction (float): Fraction of force updates
                        per epoch. Default to 0.02315.
                    short_energy_error_threshold (float): RMSE threshold for
                        energy update candidates. Default to 0.0.
                    short_force_error_threshold (float): RMSE threshold for
                        force update candidates. Default to 1.0.
                    rmse_threshold_trials (int): Maximum number of RMSE
                        threshold trials. Default to 3.
                    weights_min (float): Minimum value for initial random
                        weights. Default to -1.
                    weights_max (float): Maximum value for initial random
                        weights. Default to 1.
                    write_trainpoints (int): Write energy comparison every
                        this many epochs. Default to 1.
                    write_trainforces (int): Write force comparison every
                        this many epochs. Default to 1.
                    write_weights_epoch (int): Write weights every this many
                        epochs. Default to 1.
                    write_neuronstats (int): Write neuron statistics every
                        this many epochs. Default to 1.

                    # Kalman Filter
                    kalman_type (int): Kalman filter type. Default to 0.
                    kalman_epsilon (float): General Kalman filter parameter
                        epsilon. Default to 0.01.
                    kalman_q0 (float): General Kalman filter parameter q0.
                        Default to 0.01.
                    kalman_qtau (float): General Kalman filter parameter
                        qtau. Default to 2.302.
                    kalman_qmin (float): General Kalman filter parameter qmin.
                        Default to 1e-6.
                    kalman_eta (float): Standard Kalman filter parameter eta.
                        Default to 0.01.
                    kalman_etatau (float): Standard Kalman filter parameter
                        etatau. Defaul to 2.302.
                    kalman_etamax (float): Standard Kalman filter parameter
                        etamax. Default to 1.0.

                Symmetry functions:
                    r_cut (float): Cutoff distance (unit: Å).
                    r_etas (numpy.array): η in radial function.
                    r_shift (numpy.array): Rs in radial function.
                    a_etas (numpy.array): η in angular function.
                    zetas (numpy.array): ζ in angular function.
                    lambdas (numpy.array): λ in angular function. Default to (1, -1).
        """
        filename = 'input.nn'

        head_formatter = '{:<32s}{value}'
        type2_format = 'symfunction_short {central_atom}  2 {neighbor_atom}' \
                       '    {r_eta:.7f}    {rs:.7f}    {rcut:.7f}'
        type3_format = 'symfunction_short {central_atom}  3 {neighbor_atom1} ' \
                       '{neighbor_atom2}    {a_eta:.7f} {lambd:>2d} {zeta:.7f}   '\
                       '{rcut:.7f}'

        specie = self.specie.name
        lines = [head_formatter.format('number_of_elements', value=1),
                 head_formatter.format('elements', value=specie)]

        PARAMS = {'general': ['cutoff_type', 'scale_features', 'scale_min_short',
                              'scale_max_short', 'hidden_layers'],
                  'additional': ['epochs', 'updater_type', 'parallel_mode',
                                 'update_strategy', 'selection_mode', 'random_seed',
                                 'test_fraction', 'force_weight', 'short_energy_fraction',
                                 'short_force_fraction', 'short_energy_error_threshold',
                                 'short_force_error_threshold', 'rmse_threshold_trials',
                                 'weights_min', 'weights_max', 'write_trainpoints',
                                 'write_trainforces', 'write_weights_epoch',
                                 'write_neuronstats', 'kalman_type', 'kalman_epsilon',
                                 'kalman_q0', 'kalman_qtau', 'kalman_qmin', 'kalman_eta',
                                 'kalman_etatau', 'kalman_etamax']}
        if self.fitted:
            if self.param.get('atom_energy'):
                lines.append(head_formatter.format('atom_energy',
                                value=' '.join([specie, str(self.param.get('atom_energy'))])))
            for tag in PARAMS.get('general'):
                if tag == 'scale_features':
                    lines.append(NNinput_params.get('general').get(tag).get(self.param.get(tag)))
                elif tag == 'hidden_layers':
                    layers = self.param.get('hidden_layers')
                    activations = self.param.get('activations')
                    lines.append(head_formatter.format('global_hidden_layers_short',
                                                       value=len(layers)))
                    lines.append(head_formatter.format('global_nodes_short',
                                                       value=' '.join([str(i) for i in layers])))
                    lines.append(head_formatter.format('global_activation_short',
                                                       value=' '.join([activations] \
                                                        * len(layers) + ['l'])))
                else:
                    lines.append(head_formatter.format(tag, value=self.param.get(tag)))
            if self.normalized_nodes:
                lines.append('normalize_nodes')

            for tag in PARAMS.get('additional'):
                lines.append(head_formatter.format(tag, value=self.param.get(tag)))
            lines.append('use_short_forces')

            central_atom, neighbor_atom1, neighbor_atom2 = specie, specie, specie

            r_cut = self.param.get('r_cut')
            r_cut /= self.bohr_to_angstrom
            r_shift = np.array(self.param.get('r_shift'))
            r_shift /= self.bohr_to_angstrom

            for r_eta, rs in itertools.product(self.param.get('r_etas'), r_shift):
                lines.append(type2_format.format(central_atom=central_atom,
                                                 neighbor_atom=neighbor_atom1,
                                                 r_eta=r_eta, rs=rs, rcut=r_cut))

            for a_eta, lambd, zeta in itertools.product(self.param.get('a_etas'),
                                                        self.param.get('lambdas'), \
                                                        self.param.get('zetas')):
                lines.append(type3_format.format(central_atom=central_atom,
                                                 neighbor_atom1=neighbor_atom1,
                                                 neighbor_atom2=neighbor_atom2,
                                                 a_eta=a_eta, lambd=lambd,
                                                 zeta=zeta, rcut=r_cut))
        else:
            if kwargs.get('atom_energy'):
                lines.append(head_formatter.format('atom_energy',
                                value=' '.join([specie, str(kwargs.get('atom_energy'))])))
                self.param.update({'atom_energy': kwargs.get('atom_energy')})
            for tag in PARAMS.get('general'):
                if tag == 'scale_features':
                    value = kwargs.get(tag) if kwargs.get(tag) is not None else '1'
                    lines.append(NNinput_params.get('general').get(tag).get(value))
                    self.param.update({tag: value})
                elif tag == 'hidden_layers':
                    layers = kwargs.get(tag) if kwargs.get(tag) is not None \
                            else NNinput_params.get('general').get(tag)
                    self.param.update({tag: layers})
                    activations = kwargs.get('activations') if kwargs.get('activations') \
                            is not None else NNinput_params.get('general').get('activations')
                    self.param.update({'activations': activations})
                    lines.append(head_formatter.format('global_hidden_layers_short',
                                                       value=len(layers)))
                    lines.append(head_formatter.format('global_nodes_short',
                                                       value=' '.join([str(i) for i in layers])))
                    lines.append(head_formatter.format('global_activation_short',
                                                       value=' '.join([activations] * len(layers) \
                                                                      + ['l'])))
                else:
                    value = kwargs.get(tag) if kwargs.get(tag) is not None \
                                else NNinput_params.get('general').get(tag)
                    lines.append(head_formatter.format(tag, value=value))
                    self.param.update({tag: value})
            if kwargs.get('normalize_nodes'):
                lines.append('normalize_nodes')
                self.param.update({'normalize_nodes': True})

            for tag in PARAMS.get('additional'):
                value = kwargs.get(tag) if kwargs.get(tag) is not None \
                            else NNinput_params.get('additional').get(tag)
                lines.append(head_formatter.format(tag, value=value))
                self.param.update({tag: value})
            lines.append('use_short_forces')

            central_atom, neighbor_atom1, neighbor_atom2 = specie, specie, specie

            r_cut = kwargs.get('r_cut') if kwargs.get('r_cut') is not None \
                            else NNinput_params.get('symmetry_function').get('r_cut')
            self.param.update({'r_cut': r_cut})
            r_cut /= self.bohr_to_angstrom
            r_etas = kwargs.get('r_etas') if kwargs.get('r_etas') is not None \
                            else NNinput_params.get('symmetry_function').get('r_etas')
            self.param.update({'r_etas': r_etas})
            r_shift = kwargs.get('r_shift') if kwargs.get('r_shift') is not None \
                            else NNinput_params.get('symmetry_function').get('r_shift')
            self.param.update({'r_shift': r_shift})
            r_shift = np.array(r_shift)
            r_shift /= self.bohr_to_angstrom
            a_etas = kwargs.get('a_etas') if kwargs.get('a_etas') is not None \
                            else NNinput_params.get('symmetry_function').get('a_etas')
            self.param.update({'a_etas': a_etas})
            zetas = kwargs.get('zetas') if kwargs.get('zetas') is not None \
                            else NNinput_params.get('symmetry_function').get('zetas')
            self.param.update({'zetas': zetas})
            lambdas= kwargs.get('lambdas') if kwargs.get('lambdas') is not None \
                            else NNinput_params.get('symmetry_function').get('lambdas')
            self.param.update({'lambdas': lambdas})

            for r_eta, rs in itertools.product(r_etas, r_shift):
                lines.append(type2_format.format(central_atom=central_atom,
                                                 neighbor_atom=neighbor_atom1,
                                                 r_eta=r_eta, rs=rs, rcut=r_cut))

            for a_eta, lambd, zeta in itertools.product(a_etas, lambdas, zetas):
                lines.append(type3_format.format(central_atom=central_atom,
                                                 neighbor_atom1=neighbor_atom1,
                                                 neighbor_atom2=neighbor_atom2,
                                                 a_eta=a_eta, lambd=lambd,
                                                 zeta=zeta, rcut=r_cut))

            self.num_symm_functions = len(list(itertools.product(r_etas, r_shift))) \
                                      + len(list(itertools.product(a_etas, lambdas, zetas)))
            self.layer_sizes = [self.num_symm_functions] + self.param.get('hidden_layers')

        with open(filename, 'w') as f:
            f.write('\n'.join(lines))

        self.fitted = True

        return filename

    def load_input(self, filename='input.nn'):
        """
        Load input file from trained Neural Network Potential.

        Args:
            filename (str): The input filename.
        """
        PARAMS = {'general': ['cutoff_type', 'scale_features', 'scale_min_short',
                              'scale_max_short', 'hidden_layers'],
                  'additional': ['epochs', 'updater_type', 'parallel_mode',
                                 'update_strategy', 'selection_mode', 'random_seed',
                                 'test_fraction', 'force_weight', 'short_energy_fraction',
                                 'short_force_fraction', 'short_energy_error_threshold',
                                 'short_force_error_threshold', 'rmse_threshold_trials',
                                 'weights_min', 'weights_max', 'write_trainpoints',
                                 'write_trainforces', 'write_weights_epoch',
                                 'write_neuronstats', 'kalman_type', 'kalman_epsilon',
                                 'kalman_q0', 'kalman_qtau', 'kalman_qmin', 'kalman_eta',
                                 'kalman_etatau', 'kalman_etamax']}
        str_formatify = lambda string: float(string) if '.' in string or 'e' in string \
                            else int(string)
        param = {}
        with open(filename, 'r') as f:
            lines = f.readlines()
        df = pd.DataFrame([line.split() for line in lines if "#" not in line])
        specie = Element(np.array(df[df[0] == 'elements'][1])[0])
        self.specie = specie
        self.suffix = '{:0>3d}'.format(specie.Z)

        atom_energy = float(np.array(df[df[0] == 'atom_energy'][2])[0])
        param.update({'atom_energy': atom_energy})
        for tag in PARAMS.get('general'):
            if tag == 'scale_features':
                scale_features = '1' \
                    if len(df[df[0] == 'scale_symmetry_functions']) != 0 else 0
                param.update({'scale_features': scale_features})
            elif tag == 'hidden_layers':
                hidden_layers = [int(neuron) for neuron in np.array(df[df[0] \
                                    == 'global_nodes_short'])[0][1:] if neuron]
                param.update({'hidden_layers': hidden_layers})
                activations = np.array(df[df[0] == 'global_activation_short'])[0][1]
                param.update({'activations': activations})
            else:
                value = str_formatify(np.array(df[df[0] == tag])[0][1])
                param.update({tag: value})
        if len(df[df[0] == 'normalize_nodes']) != 0:
            param.update({'normalize_nodes': True})

        for tag in PARAMS.get('additional'):
            value = str_formatify(np.array(df[df[0] == tag])[0][1])
            param.update({tag: value})

        r_cut = np.array(df[(df[0] == 'symfunction_short') & (df[2] == '2')][6],
                         dtype=np.float)[0]
        r_cut = float('{:.1f}'.format(r_cut * units.bohr_to_angstrom))
        param.update({'r_cut': r_cut})
        r_etas = np.array(np.unique(df[(df[0] == 'symfunction_short') & (df[2] == '2')][4]),
                          dtype=np.float).tolist()
        param.update({'r_etas': r_etas})
        r_shift = np.array(np.unique(df[(df[0] == 'symfunction_short') & (df[2] == '2')][5]),
                           dtype=np.float)
        r_shift = [float('{:.1f}'.format(r * units.bohr_to_angstrom)) for r in r_shift]
        param.update({'r_shift': r_shift})
        a_etas = np.array(np.unique(df[(df[0] == 'symfunction_short') & (df[2] == '3')][5]),
                          dtype=np.float).tolist()
        param.update({'a_etas': a_etas})
        lambdas = np.array(np.unique(df[(df[0] == 'symfunction_short') & (df[2] == '3')][6]),
                           dtype=np.int).tolist()
        param.update({'lambdas': lambdas})
        zetas = np.array(np.unique(df[(df[0] == 'symfunction_short') & (df[2] == '3')][7]),
                         dtype=np.float).tolist()
        param.update({'zetas': zetas})
        self.num_symm_functions = len(list(itertools.product(r_etas, r_shift))) \
                                  + len(list(itertools.product(a_etas, lambdas, zetas)))
        self.layer_sizes = [self.num_symm_functions] + hidden_layers
        self.param = param

    def load_weights(self, weights_filename):
        """
        Load weights file of trained Neural Network Potential.

        Args
            weights_filename (str): The weights file.
        """
        with open(weights_filename) as f:
            weights_lines = f.readlines()

        weight_param = pd.DataFrame([line.split() for line in weights_lines \
                                     if "#" not in line])
        weight_param.columns = ['value', 'type', 'index', 'start_layer',
                                'start_neuron', 'end_layer', 'end_neuron']

        for layer_index in range(1, len(self.layer_sizes)):
            weights_group = weight_param[(weight_param['start_layer'] == str(layer_index - 1)) \
                                         & (weight_param['end_layer'] == str(layer_index))]

            weights = np.reshape(np.array(weights_group['value'], dtype=np.float),
                                 (self.layer_sizes[layer_index - 1],
                                  self.layer_sizes[layer_index]))
            self.weights.append(weights)

            bs_group = weight_param[(weight_param['type'] == 'b') &
                                    (weight_param['start_layer'] == str(layer_index))]
            bs = np.array(bs_group['value'], dtype=np.float)
            self.bs.append(bs)

        self.weight_param = weight_param

    def load_scaler(self, scaling_filename):
        """
        Load scaling info of trained Neural Network Potential.

        Args:
            scaling_filename (str): The scaling file.
        """
        with open(scaling_filename) as f:
            scaling_lines = f.readlines()
        scaling_param = pd.DataFrame([line.split() for line in scaling_lines \
                                      if '#' not in line])
        scaling_param.column = ['e_index', 'sf_index', 'sf_min', 'sf_max', \
                                'sf_mean', 'sf_sigma']
        self.scaling_param = scaling_param

    def read_cfgs(self, filename='output.data'):
        """
        Args:
            filename (str): The configuration file to be read.
        """
        data_pool = []
        with zopen(filename, 'rt') as f:
            lines = f.read()

        block_pattern = re.compile('begin\n(.*?)end', re.S)
        lattice_pattern = re.compile('lattice(.*?)\n')
        position_pattern = re.compile('atom(.*?)\n')
        energy_pattern = re.compile('energy(.*?)\n')

        for block in block_pattern.findall(lines):
            d = {'outputs':{}}
            lattice_str = lattice_pattern.findall(block)
            lattice = Lattice(np.array([latt.split() for latt in lattice_str],
                                        dtype=np.float) * self.bohr_to_angstrom)
            position_str = position_pattern.findall(block)
            positions = pd.DataFrame([pos.split() for pos in position_str])
            positions.columns = \
                ['x', 'y', 'z', 'specie', 'charge', 'atomic_energy', 'fx', 'fy', 'fz']
            coords = np.array(positions.loc[:, ['x', 'y', 'z']], dtype=np.float)
            coords = coords * self.bohr_to_angstrom
            species = np.array(positions['specie'])
            forces = np.array(positions.loc[:, ['fx', 'fy', 'fz']], dtype=np.float)
            forces = forces / self.eV_to_Ha / self.bohr_to_angstrom
            energy_str = energy_pattern.findall(block)[0]
            energy = float(energy_str.lstrip()) / self.eV_to_Ha
            struct = Structure(lattice=lattice, species=species, coords=coords,
                               coords_are_cartesian=True)
            d['structure'] = struct.as_dict()
            d['outputs']['energy'] = energy
            d['outputs']['forces'] = forces
            d['num_atoms'] = len(struct)

            data_pool.append(d)
        _, df = convert_docs(docs=data_pool)
        return data_pool, df

    def write_param(self):
        """
        Write optimized weights file to perform energy and force prediction.
        """
        if self.weight_param is None or self.scaling_param is None:
            raise RuntimeError("The parameters should be provided.")
        weights_filename = '.'.join(['weights', self.suffix, 'data'])
        weight_formatter = '{:>18s}{:>2s}{:>10s}{:>6s}{:>6s}{:>6s}{:>6s}'
        bias_formatter = '{:>18s}{:>2s}{:>10s}{:>6s}{:>6}'
        lines = []
        for i in range(self.weight_param.shape[0]):
            if self.weight_param.iloc[i]['type'] == 'a':
                lines.append(weight_formatter.format(*self.weight_param.iloc[i]))
            else:
                lines.append(bias_formatter.format(*self.weight_param.iloc[i]))

        with open(weights_filename, 'w') as f:
            f.writelines('\n'.join(lines))

        scaling_filename = 'scaling.data'
        scaling_formatter = '{:>4s}{:>5s}  {:>22s} {:>22s} {:>22s} {:.>22s}'
        scaling_lines = []
        for i in range(self.num_symm_functions):
            scaling_lines.append(scaling_formatter.format(*self.scaling_param.iloc[i]))
        with open(scaling_filename, 'w') as f:
            f.writelines('\n'.join(scaling_lines))

        self.write_input()

        ff_settings = [self.pair_style, self.pair_coeff.format(self.param.get('r_cut') + 1e-2)]

        return ff_settings

    def train(self, train_structures, energies=None, forces=None, stresses=None,
                                    **kwargs):
        """
        Training data with moment tensor method.

        Args:
            train_structures ([Structure]): The list of Pymatgen Structure object.
                energies ([float]): The list of total energies of each structure
                in structures list.
            energies ([float]): List of total energies of each structure in
                structures list.
            forces ([np.array]): List of (m, 3) forces array of each structure
                with m atoms in structures list. m can be varied with each
                single structure case.
            stresses (list): List of (6, ) virial stresses of each
                structure in structures list.
            kwargs: Parameters in write_input method.
        """
        if not which('nnp-train'):
            raise RuntimeError("NNP Trainer has not been found.")

        train_pool = pool_from(train_structures, energies, forces, stresses)
        atoms_filename = 'input.data'

        with ScratchDir('.'):
            atoms_filename = self.write_cfgs(filename=atoms_filename, cfg_pool=train_pool)
            output = 'training_output'

            input_filename = self.write_input(**kwargs)
            p_scaling = subprocess.Popen(['nnp-scaling', input_filename])
            stdout = p_scaling.communicate()[0]

            p_train = subprocess.Popen(['nnp-train', input_filename],
                                       stdout=open(output, 'w'))
            stdout = p_train.communicate()[0]

            rc = p_train.returncode
            if rc != 0:
                error_msg = 'RuNNer exited with return code %d' % rc
                msg = stdout.decode("utf-8").split('\n')[:-1]
                try:
                    error_line = [i for i, m in enumerate(msg)
                                  if m.startswith('ERROR')][0]
                    error_msg += ', '.join([e for e in msg[error_line:]])
                except:
                    error_msg += msg[-1]
                raise RuntimeError(error_msg)

            with zopen(output) as f:
                error_lines = f.read()

            energy_rmse_pattern = re.compile('ENERGY\s*\S*\s*(\S*)\s*(\S*).*?\n')
            forces_rmse_pattern = re.compile('FORCES\s*\S*\s*(\S*)\s*(\S*).*?\n')
            self.train_energy_rmse, self.validation_energy_rmse = \
                    np.array([line for line in energy_rmse_pattern.findall(error_lines)],
                             dtype=np.float).T
            self.train_forces_rmse, self.validation_forces_rmse = \
                    np.array([line for line in forces_rmse_pattern.findall(error_lines)],
                             dtype=np.float).T

            weights_filename_pattern = 'weights*{}.out'.format(self.param.get('epochs'))
            weights_filename = glob.glob(weights_filename_pattern)[0]

            self.suffix = weights_filename.split('.')[1]

            self.load_weights(weights_filename)
            self.load_scaler('scaling.data')

        return rc

    def evaluate(self, test_structures, ref_energies, ref_forces, ref_stresses):
        """
        Evaluate energies, forces and stresses of structures with trained
        interatomic potentials.

        Args:
            test_structures ([Structure]): List of Pymatgen Structure Objects.
            ref_energies ([float]): List of DFT-calculated total energies of
                each structure in structures list.
            ref_forces ([np.array]): List of DFT-calculated (m, 3) forces of
                each structure with m atoms in structures list. m can be varied
                with each single structure case.
            ref_stresses (list): List of DFT-calculated (6, ) viriral stresses
                of each structure in structures list.
        """
        if not which('nnp-predict'):
            raise RuntimeError("NNP Predictor has not been found.")

        original_file = 'input.data'
        predict_file = 'output.data'

        predict_pool = pool_from(test_structures, ref_energies,
                                 ref_forces, ref_stresses)
        with ScratchDir('.'):
            _, _ = self.write_param()
            original_file = self.write_cfgs(original_file, cfg_pool=predict_pool)
            _, df_orig = self.read_cfgs(original_file)

            input_filename = self.write_input()

            dfs = []
            for data in predict_pool:
                _ = self.write_cfgs(original_file, cfg_pool=[data])
                p = subprocess.Popen(['nnp-predict', input_filename], stdout=subprocess.PIPE)
                stdout = p.communicate()[0]

                rc = p.returncode
                if rc != 0:
                    error_msg = 'RuNNer exited with return code %d' % rc
                    msg = stdout.decode("utf-8").split('\n')[:-1]
                    try:
                        error_line = [i for i, m in enumerate(msg)
                                      if m.startswith('ERROR')][0]
                        error_msg += ', '.join([e for e in msg[error_line:]])
                    except:
                        error_msg += msg[-1]
                    raise RuntimeError(error_msg)

                _, df = self.read_cfgs(predict_file)
                dfs.append(df)
            df_predict = pd.concat(dfs, ignore_index=True)

        return df_orig, df_predict

    def predict(self, structure):
        """
        Predict energy, forces and stresses of the structure.

        Args:
            structure (Structure): Pymatgen Structure object.

        Returns:
            energy, forces, stress
        """
        calculator = EnergyForceStress(self)
        energy, forces, stress = calculator.calculate(structures=[structure])[0]
        return energy, forces, stress

    @staticmethod
    def from_config(input_filename, scaling_filename, weights_filename):
        """
        Initialize potentials with parameters file.

        Args:
            input_filename (str): The file storing the input configuration of
                Neural Network Potential.
            scaling_filename (str): The file storing scaling info of
                Neural Network Potential.
            weights_filename (str): The file storing weights of
                Neural Network Potential.
        """
        nnp = NNPotential()
        nnp.load_input(input_filename)
        nnp.load_scaler(scaling_filename)
        nnp.load_weights(weights_filename)
        nnp.fitted = True

        return nnp