import sympy as sp
import unittest
from sympy import I
from sympy.core.function import Function

from .expr import single_out_special_terms, product
from functools import reduce

pi = sp.S.Pi

class lambda_sqrt(Function):
    nargs = (3,)

    def _eval_expand_func(self, **hints):
        s, m_A, m_B = self.args
        return sp.sqrt((s - (m_A-m_B)**2)*(s - (m_A+m_B)**2))

    @classmethod
    def eval(cls, s, m_A, m_B):
        if s.is_zero:
            return 0

        if m_A == m_B:
            return s*sigma(s, m_A)

class sigma(Function):
    nargs = (2,)

    def _eval_expand_func(self, **hints):
        s, m = self.args
        return sp.sqrt(1-(4*m**2)/sp.expand_func(s))

    def fdiff(self, argindex=1):
        if argindex == 1:
            s, m = self.args
            return (2*m**2)/self.func(s, m)/s**2
        else:
            raise ArgumentIndexError(self, argindex)

    @classmethod
    def eval(cls, s, m):
        if s.is_zero:
            return 0

        if m.is_symbol and (m in s.free_symbols) and (m not in (4*m**2 - s).free_symbols):
            return 1/sigma(4*m**2 - s, m)

class JBar(Function):
    nargs = (2,)

    def _eval_expand_func(self, **hints):
        s, m = sp.expand_func(self.args[0]), sp.expand_func(self.args[1])
        return sp.expand_func((2 + sigma(s, m)*sp.log((sigma(s, m) - 1) / (sigma(s, m) + 1)))/(16*sp.S.Pi**2))

    def _eval_expand_into_sigmas(self, **hints):
        s, m = self.args
        return (2 + sigma(s, m)*sp.log((sigma(s, m) - 1) / (sigma(s, m) + 1)))/(16*pi**2)

    def _eval_expand_in_ks(self, **hints):
        s, m = self.args
        return (2*k_0(s, m)/pi + k_2(s, m)/2/pi)/sigma(s, m)

    @classmethod
    def eval(cls, s, m):
        if s.is_zero:
            return 0

class tau(Function):
    nargs = (2,)

    def _eval_expand_func(self, **hints):
        t, m = sp.expand_func(self.args[0]), sp.expand_func(self.args[1])
        return (sigma(t, m) - 1)/(sigma(t, m) + 1)

    def _eval_expand_into_sigmas(self, **hints):
        t, m = self.args
        return (sigma(t, m) - 1)/(sigma(t, m) + 1)

    @classmethod
    def eval(cls, t, m):
        if t.is_zero:
            return 1

def expand_into_sigmas(expr, deep=True):
    return sp.sympify(expr).expand(deep=deep, into_sigmas=True, func=False, basic=False,
    log=False, mul=False, power_exp=False, power_base=False, multinomial=False)

def expand_in_ks(expr, deep=True):
    return sp.sympify(expr).expand(deep=deep, in_ks=True, func=False, basic=False,
    log=False, mul=False, power_exp=False, power_base=False, multinomial=False)

def expand_once(expr, deep=True):
    return sp.sympify(expr).expand(deep=deep, once=True, func=False, basic=False,
    log=False, mul=False, power_exp=False, power_base=False, multinomial=False)


class kFuncApplBase(sp.Function):
    nargs = (2,)
    DUMMY_T, DUMMY_M = sp.Dummy("t"), sp.Dummy("m")

    def _eval_expand_func(self, **hints):
        t, m = self.args
        return sp.sympify(self.EXPANSION).subs({
            self.DUMMY_T: sp.expand_func(t), self.DUMMY_M: m},
            simultaneous=True)

    def _eval_expand_into_sigmas(self, **hints):
        t, m = self.args
        return sp.sympify(self.EXPANSION).subs({
            self.DUMMY_T: expand_into_sigmas(t), self.DUMMY_M: m},
            simultaneous=True)

class kFuncClass(sp.FunctionClass):
    @classmethod
    def new(mcl, order, expansion):
        _dict = dict()
        _dict["EXPANSION"] = expansion
        return mcl("k_%d" % order,
                        (kFuncApplBase,), _dict)

class IFuncApplBase(sp.Function):
    nargs = (2,)
    DUMMY_T, DUMMY_M = kFuncApplBase.DUMMY_T, kFuncApplBase.DUMMY_M

    def _eval_expand_func(self, **hints):
        t, m = self.args
        return sp.sympify(sp.expand_func(self.EXPANSION)).subs({
            self.DUMMY_T: sp.expand_func(t), self.DUMMY_M: m},
            simultaneous=True)

    def _eval_expand_once(self, **hints):
        t, m = self.args
        return sp.sympify(self.EXPANSION).subs({
            self.DUMMY_T: expand_once(t), self.DUMMY_M: m},
            simultaneous=True)

    def _eval_expand_into_sigmas(self, **hints):
        t, m = self.args
        return sp.sympify(expand_into_sigmas(expand_into_sigmas(self.EXPANSION))).subs({
            self.DUMMY_T: t, self.DUMMY_M: m},
            simultaneous=True)

    def _eval_expand_in_ks(self, **hints):
        t, m = self.args
        if self.EXPANSION_ks is None:
            raise NotImplementedError(f"no k expansion for {type(self)}")
        return sp.sympify(self.EXPANSION_ks).subs({
            self.DUMMY_T: t, self.DUMMY_M: m},
            simultaneous=True)

    @classmethod
    def eval(cls, t, m):
        if t.is_zero:
            return 0

class IFuncClass(sp.FunctionClass):
    @classmethod
    def new(mcl, order, expansion, expansion_ks=None):
        _dict = dict()
        _dict["EXPANSION"] = expansion
        _dict["EXPANSION_ks"] = expansion_ks
        return mcl("I_%d" % order,
                        (IFuncApplBase,), _dict)

t, m = IFuncApplBase.DUMMY_T, IFuncApplBase.DUMMY_M
tau_ = tau(t, m)
L_ = sp.log((1-sigma(t, m)) / (1+sigma(t, m)))

k_0 = kFuncClass.new(0, sigma(t, m)/(16*pi))
k_1 = kFuncClass.new(1, L_/(8*pi))
k_2 = kFuncClass.new(2, sigma(t, m)**2*L_/(8*pi))
k_3 = kFuncClass.new(3, (m**2/(t*sigma(t, m)))*L_**2*3/(16*pi))
k_4 = kFuncClass.new(4, (m**2/(t*sigma(t, m)))*( 1 + L_/sigma(t, m) + m**2/(t - 4*m**2)*L_**2)/(16*pi))

def normalize_k4(expr, s, m):
    ret = sp.S.Zero
    for d in single_out_special_terms(expr, s):
        if product(*d[1:]) == k_4(s, m):
            p = sp.Poly(d[0], s)
            q, r = sp.div(p, sp.Poly(s/m**2 - 4, s))
            ret += q.expr * (k_0(s, m) + k_1(s, m)/2 + k_3(s, m)/3)
            ret += r.expr * k_4(s, m)
        else:
            ret += product(*d)
    return ret

I_0  = IFuncClass.new( 0, 16*(pi**2)*JBar(t, m)*t + t + m**2*sp.log(tau(t, m))**2,
                      (2*k_0(4*m**2 - t, m)*16*pi + k_1(4*m**2 - t, m)*8*pi)/sigma(4*m**2 - t, m)*t 
                       + 16*pi*k_0(4*m**2 - t, m)/sigma(4*m**2 - t, m)*t
                       - k_3(4*m**2 - t, m)*t/sigma(4*m**2 - t, m)*(16*pi)/3)
I_1  = IFuncClass.new( 1, 8*(pi**2)*JBar(t, m)*t*(t-2*m**2) + t**2/4 + t*m**2 \
                      + m**4*sp.log(tau(t, m))**2,
                      (k_0(4*m**2 - t, m)*16*pi + k_1(4*m**2 - t, m)*4*pi)/sigma(4*m**2 - t, m)*t*t \
                      + 16*pi*k_0(4*m**2 - t, m)/sigma(4*m**2 - t, m)*t**2/4 \
                      - m**2*( (-t/m**2)*k_4(4*m**2 - t, m) )*t/sigma(4*m**2 - t, m)*(16*pi) \
                )
I_2  = IFuncClass.new( 2, 16*(pi**2)/3*JBar(t, m)*t*(t**2 - t*m**2 - 6*m**4) \
                      + (t**3)/9 + (t**2)*(m**2)/2 + 2*t*(m**4) \
                      + 2*(m**6)*sp.log(tau_)**2,
                      (-6*t*(k_1(4*m**2 - t, m)*8*pi + 2*k_0(4*m**2 - t, m)*16*pi)*(m**2*t - t**2)
                      +(9*m**2*t**2 + 2*t**3)*k_0(4*m**2 - t, m)*16*pi
                      +16*6*6*pi*t*m**4*k_4(4*m**2 - t, m)*t/m**2
                      )/18/sigma(4*m**2 - t, m)
)
I_3  = IFuncClass.new( 3, 4*(pi**2)/3*JBar(t, m)*t*(3*t**3 - 2*t**2*m**2 - 10*t*m**4 - 60*m**6) \
                      + t**4/16 + 5*t**3*m**2/18 + 5*t**2*m**4/4 + 5*t*m**6 \
                      + 5*m**8*sp.log(tau_)**2,
                      (2*k_0(4*m**2 - t, m)*16*pi + k_1(4*m**2 - t, m)*8*pi)/sigma(4*m**2 - t, m) \
                      /4/3*t*(3*t**3 - 2*t**2*m**2 - 10*t*m**4) \
                      + 80*pi*m**4*t**2*k_4(4*m**2 - t, m)/sigma(4*m**2 - t, m) \
                      + 16*pi*k_0(4*m**2 - t, m)/sigma(4*m**2 - t, m) \
                      * (t**4/16 + 5*t**3*m**2/18 + 5*t**2*m**4/4) \
)

JBar_PRIMIS = {
    0: I_0,
    1: I_1,
    2: I_2,
    3: I_3,
}

class TestPrimis(unittest.TestCase):
    def test_JBar_primis(self):
        t, m = sp.symbols("t m")

        primi_pairs = [
            (JBar(t, m),      I_0(t, m)/(16*pi**2)),
            (JBar(t, m)*t,    I_1(t, m)/(16*pi**2)),
            (JBar(t, m)*t**2, I_2(t, m)/(16*pi**2)),
            (JBar(t, m)*t**3, I_3(t, m)/(16*pi**2))
        ]

        for f_, F_ in primi_pairs:
            f = expand_into_sigmas(f_, deep=True)
            F = expand_into_sigmas(F_, deep=True)
            delta = sp.expand(sp.cancel(sp.diff(F, t) - f), func=True)
            self.assertEqual(delta, sp.S.Zero)

    def test_ks_expands(self):
        a, b = sp.Dummy('a'), sp.Dummy('b')

        for I in [I_0, I_1, I_2, I_3]:
            d = expand_into_sigmas(expand_in_ks(I(a, b))) \
                    - expand_into_sigmas(I(a, b))

            self.assertEqual(sp.simplify(sp.expand_func(d)), sp.S.Zero)

class KBarFuncApplBase(sp.Function):
    nargs = (2,)
    DUMMY_T, DUMMY_M = kFuncApplBase.DUMMY_T, kFuncApplBase.DUMMY_M

    def _eval_expand_func(self, **hints):
        t, m = self.args
        return sp.sympify(sp.expand_func(self.EXPANSION)).subs({
            self.DUMMY_T: sp.expand_func(t), self.DUMMY_M: m},
            simultaneous=True)

    def _eval_expand_into_sigmas(self, **hints):
        t, m = self.args
        return sp.sympify(expand_into_sigmas(self.EXPANSION)).subs({
            self.DUMMY_T: expand_into_sigmas(t), self.DUMMY_M: m},
            simultaneous=True)

class KBarFuncClass(sp.FunctionClass):
    @classmethod
    def new(mcl, order, expansion):
        _dict = dict()
        _dict["EXPANSION"] = expansion
        return mcl("KBar_%d" % order,
                        (KBarFuncApplBase,), _dict)

KBar_1 = KBarFuncClass.new(1, 1/16/pi**2*t/(t - 4*m**2)*(16*pi**2*JBar(t, m) - 2)**2)
KBar_2 = KBarFuncClass.new(2, (t - 4*m**2)/t*KBar_1(t, m) - 1/4/pi**2)
KBar_3 = KBarFuncClass.new(3, 1/16/pi**2*m**2/(t - 4*m**2)*(t/(t - 4*m**2)*(16*pi**2*JBar(t, m) - 2)**3 + pi**2*(16*pi**2*JBar(t, m) - 2)) - 1/32)
KBar_4 = KBarFuncClass.new(4, 1/16/pi**2*m**2/(t - 4*m**2)*(16*pi**2*JBar(t, m) - 2 + 8*pi**2*KBar_1(t, m) + 16*pi**2/3*KBar_3(t, m) + pi**2/3) \
                        - 1/32/pi**2 + 1/192)

KBarBar_0 = sp.Function("KBarBar_0")
KBarBar_2 = sp.Function("KBarBar_2")

ktilde_0 = sp.Function("ktilde_0")
ktilde_1 = sp.Function("ktilde_1")
ktilde_2 = sp.Function("ktilde_2")
ktilde_3 = sp.Function("ktilde_3")

KTilde_0 = sp.Function("KTilde_0")
KTildeBar_0 = sp.Function("KTildeBar_0")
KTilde_1 = sp.Function("KTilde_1")
KTildeBar_1 = sp.Function("KTildeBar_1")
KTilde_2 = sp.Function("KTilde_2")
KTildeBar_2 = sp.Function("KTildeBar_2")
KTilde_3 = sp.Function("KTilde_3")
KTildeBar_3 = sp.Function("KTildeBar_3")

KTilde_lambd_0 = sp.Function("KLambda_0")
KTilde_lambd_1 = sp.Function("KLambda_1")
KTilde_lambd_2 = sp.Function("KLambda_2")
KTilde_lambd_3 = sp.Function("KLambda_3")

del t, m

def disp_integral(disc, var, prefix="j"):
    '''
    Evaluates a dispersion integral given an imaginary discontinuity
    in *disc*.

    That is, the returned expression F will be such that

        2*i*(*disc*) = Disc F(x)

    The function assumes *disc* contains a Heaviside function
    describing the discontinuity threshold.
    '''

    wild_m2 = sp.Wild('m2')
    m = sp.Dummy('m')
    patterns = {
        sigma(var, m):     (-1, JBar(var, m)*16*sp.S.Pi),
        k_0(var, m):       (-1, JBar(var, m), KBarBar_0(var, m)),
        k_1(var, m):       (-1, KBar_1(var, m)),
        k_2(var, m):       (-1, KBar_2(var, m), KBarBar_2(var, m)),
        k_3(var, m):       (-1, KBar_3(var, m)),
        k_4(var, m):       (-1, KBar_4(var, m)),
        ktilde_0(var, m):  (-1, KTilde_0(var, m), KTildeBar_0(var, m)),
        ktilde_1(var, m):  (-1, KTilde_1(var, m), KTildeBar_1(var, m)),
        ktilde_2(var, m, wild_m2):  (-1, KTilde_2(var, m, wild_m2), KTildeBar_2(var, m, wild_m2)),
        ktilde_3(var, m, wild_m2):  (-1, KTilde_3(var, m, wild_m2), KTildeBar_3(var, m, wild_m2)),
        ktilde_0(var, m)/lambda_sqrt(var, wild_m2, m)**2:          (-1, KTilde_lambd_0(var, m, wild_m2)),
        ktilde_1(var, m)/lambda_sqrt(var, wild_m2, m)**2:          (-1, KTilde_lambd_1(var, m, wild_m2)),
        ktilde_2(var, m, wild_m2)/lambda_sqrt(var, wild_m2, m)**2: (-1, KTilde_lambd_2(var, m, wild_m2)),
        ktilde_3(var, m, wild_m2)/lambda_sqrt(var, wild_m2, m)**2: (-1, KTilde_lambd_3(var, m, wild_m2)),
    }

    ret = sp.S.Zero

    for decomp in single_out_special_terms(sp.cancel(disc), var):
        poly, specials = decomp[0], decomp[1:]
        specials = list(specials)
        to_match = reduce(lambda a, b: a*b, specials)

        # first match a heaviside function
        match = None
        wild_m = sp.Wild('m')
        for i, special in enumerate(specials):
            match = special.match(sp.Heaviside(var-4*wild_m**2))
            if match is not None:
                specials.pop(i)
                break
        if not match:
            raise NotImplementedError("unimplemented dispersion integrand "
                                      "containing %r (1)" % to_match)

        m_found = match[wild_m]
        remainder = sp.S.One
        var_powers = 0
        sigma_powers = 0
        for special in specials:
            if special == sigma(var, m_found):
                sigma_powers += 1
            elif type(special) is sp.Pow and special.args[0] == var:
                var_powers += special.args[1]
            else:
                remainder *= special

        while sigma_powers >= 2:
            poly *= sp.expand_func(sigma(var, m_found)**2)*var
            var_powers -= 1
            sigma_powers -= 2
        if sigma_powers == 1:
            remainder *= sigma(var, m_found)
        elif sigma_powers == 0:
            pass
        else:
            raise NotImplementedError("unimplemented dispersion integrand "
                                      "containing %r (2)" % to_match)

        matched = False
        for pattern, rhs in patterns.items():
            sublimit, fs = rhs[0], rhs[1:]
            match = remainder.match(pattern.subs({m: m_found}))
            if match is None:
                continue
            matched = True

            ppoly = sp.cancel(poly).as_poly(var)
            for order, coeff in enumerate(reversed(ppoly.all_coeffs())):
                if coeff == sp.S.Zero:
                    continue
                power = order + var_powers
                if power < sublimit:
                    overrun = sublimit - power
                    if overrun >= len(fs):
                        raise NotImplementedError("unimplemented higher "
                                      "subtracted integral of %r (overrun %d, coeff %r)" \
                                      % (pattern, overrun, coeff))
                    f = fs[overrun]
                else:
                    f = fs[0]
                ret += f.subs({m: m_found, wild_m2: match.get(wild_m2, None)}) \
                        * coeff * var**power

        if not matched:
            raise NotImplementedError("unimplemented dispersion integrand "
                                      "containing %r (3)" % to_match)

    return ret
