def test_chi_squared():
    k = Symbol("k", integer=True)

    X = ChiSquared('x', k)
    assert density(X)(x) == 2**(-k/2)*x**(k/2 - 1)*exp(-x/2)/gamma(k/2)
    assert cdf(X)(x) == Piecewise((lowergamma(k/2, x/2)/gamma(k/2), x >= 0), (0, True))

    X = ChiSquared('x', 15)
    assert cdf(X)(3) == -14873*sqrt(6)*exp(-S(3)/2)/(5005*sqrt(pi)) + erf(sqrt(6)/2)

def test_cauchy():
    x0 = Symbol("x0")
    gamma = Symbol("gamma", positive=True)

    X = Cauchy('x', x0, gamma)
    assert density(X)(x) == 1/(pi*gamma*(1 + (x - x0)**2/gamma**2))
    assert cdf(X)(x) == atan((x - x0)/gamma)/pi + S.Half
    assert diff(cdf(X)(x), x) == density(X)(x)

    gamma = Symbol("gamma", positive=False)
    raises(ValueError, lambda: Cauchy('x', x0, gamma))

def test_logistic():
    mu = Symbol("mu", real=True)
    s = Symbol("s", positive=True)

    X = Logistic('x', mu, s)
    assert density(X)(x) == exp((-x + mu)/s)/(s*(exp((-x + mu)/s) + 1)**2)
    assert cdf(X)(x) == 1/(exp((mu - x)/s) + 1)

def symSq():
    rho = r / sigma - 0.5 * sigma
    N = stats.Normal('N', 0, 1)
    theta2t = 1/(2*T)*theta2/sqrt(1+theta2**2)
    # theta2t = 1/(2*T)*(exp(theta2)-1)/(exp(theta2)+1)
    a = (1/(2*T)+theta2t) * 4*sqrt(T)/sigma*log(B/s0) - theta1
    Kprime = K * exp(-r*T)
    g1 = stats.cdf(N)(a - log(Kprime/s0)/sigma/sqrt(T) + 3/2*sigma*sqrt(T))
    g2 = stats.cdf(N)(a - log(Kprime/s0)/sigma/sqrt(T) + 0.5*sigma*sqrt(T))
    g3 = stats.cdf(N)(a - log(Kprime/s0)/sigma/sqrt(T) - 0.5*sigma*sqrt(T))
    pre = exp((1/(2*T)+theta2t)*(4*T*rho/sigma*log(B/s0) - 4/sigma**2*log(B/s0)**2))
    T1 = s0**2*g1*exp(0.5*(a+2*sigma*sqrt(T))**2 - sigma**2*T)
    T2 = s0*Kprime*g2*exp(0.5*(a+sigma*sqrt(T))**2 - 0.5*sigma**2*T)
    T3 = Kprime**2*g3*exp(0.5*a**2)
    res = pre * (T1 - 2*T2 + T3) * exp(0.5*theta1**2) / (1-4*T**2*theta2t**2)
    return res

def test_dagum():
    p = Symbol("p", positive=True)
    b = Symbol("b", positive=True)
    a = Symbol("a", positive=True)

    X = Dagum('x', p, a, b)
    assert density(X)(x) == a*p*(x/b)**(a*p)*((x/b)**a + 1)**(-p - 1)/x
    assert cdf(X)(x) == Piecewise(((1 + (x/b)**(-a))**(-p), x >= 0),
                                    (0, True))

def test_cdf():
    X = Normal('x', 0, 1)

    d = cdf(X)
    assert P(X < 1) == d(1)
    assert d(0) == S.Half

    d = cdf(X, X > 0)  # given X>0
    assert d(0) == 0

    Y = Exponential('y', 10)
    d = cdf(Y)
    assert d(-5) == 0
    assert P(Y > 3) == 1 - d(3)

    raises(ValueError, lambda: cdf(X + Y))

    Z = Exponential('z', 1)
    f = cdf(Z)
    z = Symbol('z')
    assert f(z) == Piecewise((1 - exp(-z), z >= 0), (0, True))

def test_gamma():
    k = Symbol("k", positive=True)
    theta = Symbol("theta", positive=True)

    X = Gamma('x', k, theta)
    assert density(X) == Lambda(_x,
                                _x**(k - 1)*theta**(-k)*exp(-_x/theta)/gamma(k))
    assert cdf(X, meijerg=True) == Lambda(_z, Piecewise(
                                          (-k*lowergamma(k, 0)/gamma(k + 1) + k*lowergamma(k, _z/theta)/gamma(k + 1), _z >= 0), (0, True)))
    assert variance(X) == (-theta**2*gamma(k + 1)**2/gamma(k)**2 +
           theta*theta**(-k)*theta**(k + 1)*gamma(k + 2)/gamma(k))

    k, theta = symbols('k theta', real=True, bounded=True, positive=True)
    X = Gamma('x', k, theta)
    assert simplify(E(X)) == k*theta
    # can't get things to simplify on this one so we use subs
    assert variance(X).subs(k, 5) == (k*theta**2).subs(k, 5)

def test_prob():
    def emit(name, iname, cdf, args, no_small=False):
        V = []
        for arg in sorted(args):
            y = cdf(*arg)
            if isinstance(y, mpf):
                e = sp.nsimplify(y, rational=True)
                if e.is_Rational and e.q <= 1000 and \
                        mp.almosteq(mp.mpf(e), y, 1e-25):
                    y = e
            else:
                y = N(y)
            V.append(arg + (y,))
        for v in V:
            if name:
                test(name, *v)
        for v in V:
            if iname and (not no_small or 1/1000 <= v[-1] <= 999/1000):
                test(iname, *(v[:-2] + v[:-3:-1]))

    x = sp.Symbol("x")
    emit("ncdf", "nicdf",
         sp.Lambda(x, st.cdf(st.Normal("X", 0, 1))(x)), zip(exparg))
    # using cdf() for anything more complex is too slow

    df = FiniteSet(1, S(3)/2, 2, S(5)/2, 5, 25)
    emit("c2cdf", "c2icdf",
         lambda k, x: sp.lowergamma(k/2, x/2)/sp.gamma(k/2),
         ProductSet(df, posarg), no_small=True)

    dfint = df & sp.fancysets.Naturals()
    def cdf(k, x):
        k, x = map(mpf, (k, x))
        return .5 + .5*mp.sign(x)*mp.betainc(k/2, .5, x1=1/(1+x**2/k),
                                             regularized=True)
    emit("stcdf", "sticdf", cdf, ProductSet(dfint, exparg))

    def cdf(d1, d2, x):
        d1, d2, x = map(mpf, (d1, d2, x))
        return mp.betainc(d1/2, d2/2, x2=x/(x+d2/d1), regularized=True)

    emit("fcdf", "ficdf", cdf, ProductSet(dfint, dfint, posarg))

    kth = ProductSet(sp.ImageSet(lambda x: x/5, df),
                     posarg - FiniteSet(0))
    emit("gcdf", "gicdf",
         lambda k, th, x: sp.lowergamma(k, x/th)/sp.gamma(k),
         ProductSet(kth, posarg), no_small=True)

    karg = FiniteSet(0, 1, 2, 5, 10, 15, 40)
    knparg = [(k, n, p) for k, n, p
              in ProductSet(karg, karg, posarg & Interval(0, 1, True, True))
              if k <= n and n > 0]
    def cdf(k, n, p):
        return st.P(st.Binomial("X", n, p) <= k)
    emit("bncdf", "bnicdf", cdf, knparg, no_small=True)

    def cdf(k, lamda):
        return sp.uppergamma(k+1, lamda)/sp.gamma(k+1)
    emit("pscdf", "psicdf", cdf,
         ProductSet(karg, posarg + karg - FiniteSet(0)), no_small=True)

    x, i = sp.symbols("x i")
    def smcdf(n, e):
        return 1-sp.Sum(sp.binomial(n, i)*e*(e+i/n)**(i-1)*(1-e-i/n)**(n-i),
                        (i, 0, sp.floor(n*(1-e)))).doit()
    kcdf = sp.Lambda(x,
        sp.sqrt(2*pi)/x*sp.Sum(sp.exp(-pi**2/8*(2*i-1)**2/x**2), (i, 1, oo)))
    smarg = ProductSet(karg - FiniteSet(0), posarg & Interval(0, 1, True, True))
    karg = FiniteSet(S(1)/100, S(1)/10) + (posarg & Interval(S(1)/4, oo, True))

    for n, e in sorted(smarg):
        test("smcdf", n, e, N(smcdf(n, e)))
    prec("1e-10")
    for x in sorted(karg):
        test("kcdf", x, N(kcdf(x)))
    prec("1e-9")
    for n, e in sorted(smarg):
        p = smcdf(n, e)
        if p < S(9)/10:
            test("smicdf", n, N(p), e)
    prec("1e-6")
    for x in sorted(karg):
        p = kcdf(x)
        if N(p) > S(10)**-8:
            test("kicdf", N(p), x)

def test_cdf():
    D = Die('D', 6)
    o = S.One

    assert cdf(
        D) == sympify({1: o/6, 2: o/3, 3: o/2, 4: 2*o/3, 5: 5*o/6, 6: o})

"""
Created on Thu Mar 28 15:57:09 2019

@author: javie
"""


# SymPy Imports
from sympy import init_session
init_session()


gini = 0.7
pLambda = (1 + gini) / (2 * gini)
population = 4000
thetaMin = 0.1
cParam = 0.01
gamma = 2
delta = 0.5
cFijo = 0
discountQ = 0.7
probRichest = 0.99

p, q, th = symbols('p q th', real=True)
dem = symbols('dem', integer=True)

from sympy.stats import P, E, variance, Binomial, cdf, density, Pareto

pa = Pareto("pa", thetaMin, pLambda)
bi = Binomial("bi", population, 1-cdf(pa)(th))