def test_pspace():
    X, Y = Normal(0,1), Normal(0,1)

    assert not pspace(5+3)
    assert pspace(X) == X.pspace
    assert pspace(2*X+1) == X.pspace
    assert pspace(2*X+Y) == ProductPSpace(Y.pspace, X.pspace)

def test_domains():
    X, Y = Die('x', 6), Die('y', 6)
    x, y = X.symbol, Y.symbol
    # Domains
    d = where(X > Y)
    assert d.condition == (x > y)
    d = where(And(X > Y, Y > 3))
    assert d.as_boolean() == Or(And(Eq(x, 5), Eq(y, 4)), And(Eq(x, 6),
        Eq(y, 5)), And(Eq(x, 6), Eq(y, 4)))
    assert len(d.elements) == 3

    assert len(pspace(X + Y).domain.elements) == 36

    Z = Die('x', 4)

    raises(ValueError, lambda: P(X > Z))  # Two domains with same internal symbol

    pspace(X + Y).domain.set == FiniteSet(1, 2, 3, 4, 5, 6)**2

    assert where(X > 3).set == FiniteSet(4, 5, 6)
    assert X.pspace.domain.dict == FiniteSet(
        *[Dict({X.symbol: i}) for i in range(1, 7)])

    assert where(X > Y).dict == FiniteSet(*[Dict({X.symbol: i, Y.symbol: j})
            for i in range(1, 7) for j in range(1, 7) if i > j])

def test_ProductPSpace():
    X = Normal('X', 0, 1)
    Y = Normal('Y', 0, 1)
    px = X.pspace
    py = Y.pspace
    assert pspace(X + Y) == ProductPSpace(px, py)
    assert pspace(X + Y) == ProductPSpace(py, px)

def test_domains():
    x, y = symbols('x y')
    X, Y= Die(6, symbol=x), Die(6, symbol=y)
    # Domains
    d = Where(X>Y)
    assert d.condition == (x > y)
    d = Where(And(X>Y, Y>3))
    assert d.as_boolean() == Or(And(Eq(x,5), Eq(y,4)), And(Eq(x,6), Eq(y,5)),
        And(Eq(x,6), Eq(y,4)))
    assert len(d.elements) == 3

    assert len(pspace(X+Y).domain.elements) == 36

    Z = Die(4, symbol=x)

    raises(ValueError, "P(X>Z)") # Two domains with same internal symbol

    pspace(X+Y).domain.set == FiniteSet(1,2,3,4,5,6)**2

    assert Where(X>3).set == FiniteSet(4,5,6)
    assert X.pspace.domain.dict == FiniteSet(
            Dict({X.symbol:i}) for i in range(1,7))

    assert Where(X>Y).dict == FiniteSet(Dict({X.symbol:i, Y.symbol:j})
            for i in range(1,7) for j in range(1,7) if i>j)

def test_pspace():
    X, Y = Normal('X', 0, 1), Normal('Y', 0, 1)

    raises(ValueError, lambda: pspace(5 + 3))
    raises(ValueError, lambda: pspace(x < 1))
    assert pspace(X) == X.pspace
    assert pspace(2*X + 1) == X.pspace
    assert pspace(2*X + Y) == ProductPSpace(Y.pspace, X.pspace)

def test_FiniteRV():
    F = FiniteRV("F", {1: S.Half, 2: S.One / 4, 3: S.One / 4})

    assert dict(density(F).items()) == {S(1): S.Half, S(2): S.One / 4, S(3): S.One / 4}
    assert P(F >= 2) == S.Half

    assert pspace(F).domain.as_boolean() == Or(*[Eq(F.symbol, i) for i in [1, 2, 3]])

def test_RandomDomain():
    from sympy.stats import Normal, Die, Exponential, pspace, where
    X = Normal('x1', 0, 1)
    assert str(where(X > 0)) == "Domain: 0 < x1"

    D = Die('d1', 6)
    assert str(where(D > 4)) == "Domain: Or(d1 == 5, d1 == 6)"

    A = Exponential('a', 1)
    B = Exponential('b', 1)
    assert str(pspace(Tuple(A, B)).domain) == "Domain: And(0 <= a, 0 <= b)"

def test_latex_RandomDomain():
    from sympy.stats import Normal, Die, Exponential, pspace, where
    X = Normal(0, 1, symbol=Symbol('x1'))
    assert latex(where(X>0)) == "Domain: 0 < x_{1}"

    D = Die(6, symbol=Symbol('d1'))
    assert latex(where(D>4)) == r"Domain: d_{1} = 5 \vee d_{1} = 6"

    A = Exponential(1, symbol=Symbol('a'))
    B = Exponential(1, symbol=Symbol('b'))
    assert latex(pspace(Tuple(A,B)).domain) =="Domain: 0 \leq a \wedge 0 \leq b"

def test_RandomDomain():
    from sympy.stats import Normal, Die, Exponential, pspace, Where
    X = Normal(0, 1, symbol=Symbol('x1'))
    assert str(Where(X>0)) == "Domain: 0 < x1"

    D = Die(6, symbol=Symbol('d1'))
    assert str(Where(D>4)) == "Domain: Or(d1 == 5, d1 == 6)"

    A = Exponential(1, symbol=Symbol('a'))
    B = Exponential(1, symbol=Symbol('b'))
    assert str(pspace(Tuple(A,B)).domain) =="Domain: And(0 <= a, 0 <= b)"

def test_FiniteRV():
    F = FiniteRV('F', {1: S.Half, 2: S.One/4, 3: S.One/4})

    assert dict(density(F).items()) == {S(1): S.Half, S(2): S.One/4, S(3): S.One/4}
    assert P(F >= 2) == S.Half

    assert pspace(F).domain.as_boolean() == Or(
        *[Eq(F.symbol, i) for i in [1, 2, 3]])

    raises(ValueError, lambda: FiniteRV('F', {1: S.Half, 2: S.Half, 3: S.Half}))
    raises(ValueError, lambda: FiniteRV('F', {1: S.Half, 2: S(-1)/2, 3: S.One}))
    raises(ValueError, lambda: FiniteRV('F', {1: S.One, 2: S(3)/2, 3: S.Zero, 4: S(-1)/2, 5: S(-3)/4, 6: S(-1)/4}))

def test_latex_RandomDomain():
    from sympy.stats import Normal, Die, Exponential, pspace, where
    X = Normal('x1', 0, 1)
    assert latex(where(X > 0)) == "Domain: x_{1} > 0"

    D = Die('d1', 6)
    assert latex(where(D > 4)) == r"Domain: d_{1} = 5 \vee d_{1} = 6"

    A = Exponential('a', 1)
    B = Exponential('b', 1)
    assert latex(
        pspace(Tuple(A, B)).domain) == "Domain: a \geq 0 \wedge b \geq 0"

def test_latex_RandomDomain():
    from sympy.stats import Normal, Die, Exponential, pspace, where

    X = Normal("x1", 0, 1)
    assert latex(where(X > 0)) == "Domain: 0 < x_{1}"

    D = Die("d1", 6)
    assert latex(where(D > 4)) == r"Domain: d_{1} = 5 \vee d_{1} = 6"

    A = Exponential("a", 1)
    B = Exponential("b", 1)
    assert latex(pspace(Tuple(A, B)).domain) == "Domain: 0 \leq a \wedge 0 \leq b"

def test_coins():
    C, D = Coin("C"), Coin("D")
    H, T = symbols("H, T")
    assert P(Eq(C, D)) == S.Half
    assert density(Tuple(C, D)) == {(H, H): S.One / 4, (H, T): S.One / 4, (T, H): S.One / 4, (T, T): S.One / 4}
    assert dict(density(C).items()) == {H: S.Half, T: S.Half}

    F = Coin("F", S.One / 10)
    assert P(Eq(F, H)) == S(1) / 10

    d = pspace(C).domain

    assert d.as_boolean() == Or(Eq(C.symbol, H), Eq(C.symbol, T))

    raises(ValueError, lambda: P(C > D))  # Can't intelligently compare H to T

def test_coins():
    C, D = Coin(), Coin()
    H, T = sorted(density(C).keys())
    assert P(Eq(C, D)) == S.Half
    assert density(Tuple(C, D)) == {(H, H): S.One / 4, (H, T): S.One / 4, (T, H): S.One / 4, (T, T): S.One / 4}
    assert density(C) == {H: S.Half, T: S.Half}

    E = Coin(S.One / 10)
    assert P(Eq(E, H)) == S(1) / 10

    d = pspace(C).domain

    assert d.as_boolean() == Or(Eq(C.symbol, H), Eq(C.symbol, T))

    raises(ValueError, lambda: P(C > D))  # Can't intelligently compare H to T

def test_dice():
    # TODO: Make iid method!
    X, Y, Z = Die('X', 6), Die('Y', 6), Die('Z', 6)
    a, b = symbols('a b')

    assert E(X) == 3 + S.Half
    assert variance(X) == S(35)/12
    assert E(X + Y) == 7
    assert E(X + X) == 7
    assert E(a*X + b) == a*E(X) + b
    assert variance(X + Y) == variance(X) + variance(Y) == cmoment(X + Y, 2)
    assert variance(X + X) == 4 * variance(X) == cmoment(X + X, 2)
    assert cmoment(X, 0) == 1
    assert cmoment(4*X, 3) == 64*cmoment(X, 3)
    assert covariance(X, Y) == S.Zero
    assert covariance(X, X + Y) == variance(X)
    assert density(Eq(cos(X*S.Pi), 1))[True] == S.Half
    assert correlation(X, Y) == 0
    assert correlation(X, Y) == correlation(Y, X)
    assert smoment(X + Y, 3) == skewness(X + Y)
    assert smoment(X, 0) == 1
    assert P(X > 3) == S.Half
    assert P(2*X > 6) == S.Half
    assert P(X > Y) == S(5)/12
    assert P(Eq(X, Y)) == P(Eq(X, 1))

    assert E(X, X > 3) == 5 == moment(X, 1, 0, X > 3)
    assert E(X, Y > 3) == E(X) == moment(X, 1, 0, Y > 3)
    assert E(X + Y, Eq(X, Y)) == E(2*X)
    assert moment(X, 0) == 1
    assert moment(5*X, 2) == 25*moment(X, 2)

    assert P(X > 3, X > 3) == S.One
    assert P(X > Y, Eq(Y, 6)) == S.Zero
    assert P(Eq(X + Y, 12)) == S.One/36
    assert P(Eq(X + Y, 12), Eq(X, 6)) == S.One/6

    assert density(X + Y) == density(Y + Z) != density(X + X)
    d = density(2*X + Y**Z)
    assert d[S(22)] == S.One/108 and d[S(4100)] == S.One/216 and S(3130) not in d

    assert pspace(X).domain.as_boolean() == Or(
        *[Eq(X.symbol, i) for i in [1, 2, 3, 4, 5, 6]])

    assert where(X > 3).set == FiniteSet(4, 5, 6)

def test_beta():
    a, b = symbols('alpha beta', positive=True)

    B = Beta(a, b)

    assert pspace(B).domain.set == Interval(0, 1)

    x, dens = Density(B)
    assert dens == x**(a-1)*(1-x)**(b-1) / beta(a,b)

    # This is too slow
    # assert E(B) == a / (a + b)
    # assert Var(B) == (a*b) / ((a+b)**2 * (a+b+1))

    # Full symbolic solution is too much, test with numeric version
    a, b = 1, 2
    B = Beta(a, b)
    assert E(B) == a / S(a + b)
    assert Var(B) == (a*b) / S((a+b)**2 * (a+b+1))

def test_dice():
    X, Y, Z= Die(6), Die(6), Die(6)
    a,b = symbols('a b')

    assert E(X) == 3+S.Half
    assert Var(X) == S(35)/12
    assert E(X+Y) == 7
    assert E(X+X) == 7
    assert E(a*X+b) == a*E(X)+b
    assert Var(X+Y) == Var(X) + Var(Y)
    assert Var(X+X) == 4 * Var(X)
    assert Covar(X,Y) == S.Zero
    assert Covar(X, X+Y) == Var(X)
    assert Density(Eq(cos(X*S.Pi),1))[True] == S.Half

    assert P(X>3) == S.Half
    assert P(2*X > 6) == S.Half
    assert P(X>Y) == S(5)/12
    assert P(Eq(X,Y)) == P(Eq(X,1))

    assert E(X, X>3) == 5
    assert E(X, Y>3) == E(X)
    assert E(X+Y, Eq(X,Y)) == E(2*X)
    assert E(X+Y-Z, 2*X>Y+1) == S(49)/12

    assert P(X>3, X>3) == S.One
    assert P(X>Y, Eq(Y, 6)) == S.Zero
    assert P(Eq(X+Y, 12)) == S.One/36
    assert P(Eq(X+Y, 12), Eq(X, 6)) == S.One/6

    assert Density(X+Y) == Density(Y+Z) != Density(X+X)
    d = Density(2*X+Y**Z)
    assert d[S(22)] == S.One/108 and d[S(4100)]==S.One/216 and S(3130) not in d

    assert pspace(X).domain.as_boolean() == Or(
            *[Eq(X.symbol, i) for i in [1,2,3,4,5,6]])

    assert Where(X>3).set == FiniteSet(4,5,6)

def test_random_parameters():
    mu = Normal('mu', 2, 3)
    meas = Normal('T', mu, 1)
    assert density(meas, evaluate=False)(z)
    assert isinstance(pspace(meas), ProductPSpace)

def test_FinitePSpace():
    X = Die('X', 6)
    space = pspace(X)
    assert space.density == DieDistribution(6)