Repository: sympy/sympy

Issue:
Problem with discrete sampling
The code 
```python
from sympy import S
from sympy.stats import Geometric, sample
X = Geometric("X",S(1)/2)
sample(X)
```
throws an error (`global name 'floor' is not defined`). I found the same issue with `Poisson` as well. 

Discussion:

Patch:
diff --git a/sympy/stats/drv.py b/sympy/stats/drv.py
index d7c8e0a3ef..de37192599 100644
--- a/sympy/stats/drv.py
+++ b/sympy/stats/drv.py
@@ -4,6 +4,7 @@
         Piecewise, S, cacheit, Sum, exp, I, oo)
 from sympy.solvers.solveset import solveset
 from sympy.stats.rv import NamedArgsMixin, SinglePSpace, SingleDomain
+from sympy.functions.elementary.integers import floor
 import random
 
 class SingleDiscreteDistribution(Basic, NamedArgsMixin):
@@ -30,7 +31,10 @@ def check(*args):
     def sample(self):
         """ A random realization from the distribution """
         icdf = self._inverse_cdf_expression()
-        return floor(icdf(random.uniform(0, 1)))
+        while True:
+            sample_ = floor(list(icdf(random.uniform(0, 1)))[0])
+            if sample_ >= self.set.inf:
+                return sample_
 
     @cacheit
     def _inverse_cdf_expression(self):
@@ -38,16 +42,18 @@ def _inverse_cdf_expression(self):
 
         Used by sample
         """
-        x, z = symbols('x, z', real=True, positive=True, cls=Dummy)
+        x = symbols('x', positive=True,
+         integer=True, cls=Dummy)
+        z = symbols('z', positive=True, cls=Dummy)
+        cdf_temp = self.cdf(x)
         # Invert CDF
         try:
-            inverse_cdf = list(solveset(self.cdf(x) - z, x))
+            inverse_cdf = solveset(cdf_temp - z, x, domain=S.Reals)
         except NotImplementedError:
             inverse_cdf = None
-        if not inverse_cdf or len(inverse_cdf) != 1:
+        if not inverse_cdf or len(inverse_cdf.free_symbols) != 1:
             raise NotImplementedError("Could not invert CDF")
-
-        return Lambda(z, inverse_cdf[0])
+        return Lambda(z, inverse_cdf)
 
     @cacheit
     def compute_cdf(self, **kwargs):
diff --git a/sympy/stats/drv_types.py b/sympy/stats/drv_types.py
index 5b14f5d953..9f819b3ae8 100644
--- a/sympy/stats/drv_types.py
+++ b/sympy/stats/drv_types.py
@@ -3,6 +3,8 @@
 from sympy.stats.drv import SingleDiscreteDistribution, SingleDiscretePSpace
 from sympy import factorial, exp, S, sympify
 from sympy.stats.rv import _value_check
+from sympy.sets.sets import Interval
+import random
 
 __all__ = ['Geometric', 'Poisson']
 
@@ -26,6 +28,13 @@ def check(lamda):
     def pdf(self, k):
         return self.lamda**k / factorial(k) * exp(-self.lamda)
 
+    def sample(self):
+        u = random.uniform(0, 1)
+        inf = self.set.inf
+        while u > self.cdf(inf):
+            inf += 1
+        return inf
+
 
 def Poisson(name, lamda):
     r"""


Test patch:
diff --git a/sympy/stats/tests/test_discrete_rv.py b/sympy/stats/tests/test_discrete_rv.py
index 8965da92f1..2604d4232d 100644
--- a/sympy/stats/tests/test_discrete_rv.py
+++ b/sympy/stats/tests/test_discrete_rv.py
@@ -1,8 +1,9 @@
 from sympy.stats.drv_types import (PoissonDistribution, GeometricDistribution,
-        Poisson)
+        Poisson, Geometric)
 from sympy.abc import x
 from sympy import S, Sum
 from sympy.stats import E, variance, density, characteristic_function
+from sympy.stats.rv import sample
 
 def test_PoissonDistribution():
     l = 3
@@ -28,3 +29,8 @@ def test_GeometricDistribution():
     assert d.expectation(x**2, x) - d.expectation(x, x)**2 == (1-p)/p**2
     assert abs(d.cdf(20000).evalf() - 1) < .001
     assert d.characteristic_function(0).doit() == 1
+
+def test_sample():
+    X, Y = Geometric('X', S(1)/2), Poisson('Y', 4)
+    assert sample(X) in X.pspace.domain.set
+    assert sample(Y) in Y.pspace.domain.set