diff --git a/06_krisis.py b/06_krisis.py
index 07250cc..3a8f0bc 100644
--- a/06_krisis.py
+++ b/06_krisis.py
@@ -13,31 +13,8 @@ def from_angles_1(theta, phi):
     return State(m.m)
 
 
-class MemoizedExp(object):
-    def __init__(self):
-        self.theta = None
-        self.phi = None
-
-    def gen_exp_for_cirq_0(self):
-        if not (self.theta and self.phi):
-            self.theta = round(np.random.uniform(0, 1), 10)
-            self.phi = round(np.random.uniform(0, 1), 10)
-        return self.theta, self.phi
-
-    def reset(self):
-        self.theta = None
-        self.phi = None
-
-
-def gen_exp_for_cirq_1():
-    """TODO: How to generate the exponents for the second case"""
-    theta = round(np.random.uniform(0, 1), 10)
-    phi = round(np.random.uniform(0, 1), 10)
-    return theta, phi
-
-
-def print_all_samples(all_samples):
-    print("------------- ALL SAMPLES for cirq")
+def print_all_samples(name, all_samples):
+    print("------------- ALL SAMPLES for {}".format(name))
     for k, v in sorted(all_samples.items(), key=lambda x: x[0]):
         print("{}: {}".format(k, v))
     print("==============================================")
@@ -62,14 +39,48 @@ def math_sim(q2func=from_angles_1, iterations=1000, sample_count=1):
         this_samples = qp.get_sample(sample_count)
         for k, v in this_samples.items():
             all_samples[k] += v
-    print_all_samples(all_samples)
+    print_all_samples("math_sim", all_samples)
 
 
-def cirq_sim(q1func, q2func, memoized_exp, iterations=1000, sample_count=1):
+class MemoizedExp(object):
+    def __init__(self, *funcs):
+        self.memoized = dict()
+        self.funcs = {func.__name__: func for func in funcs}
+        self.reset()
+
+    def call(self, func_name, *args, **kwargs):
+        if func_name not in self.funcs:
+            raise Exception("Func {} doesn't exist".format(func_name))
+        if self.memoized[func_name]:
+            return self.memoized[func_name]
+        values = self.funcs[func_name](*args, **kwargs)
+        self.memoized[func_name] = values
+        return values
+
+    def reset(self):
+        for func_name in self.funcs.keys():
+            self.memoized[func_name] = None
+
+
+def gen_exp_for_cirq_0():
+    theta = round(np.random.uniform(0, 1), 10)
+    phi = round(np.random.uniform(0, 1), 10)
+    return theta, phi
+
+
+def gen_exp_for_cirq_1():
+    """TODO: How to generate the exponents for the second case"""
+    theta = round(np.random.uniform(0, 1), 10)
+    phi = round(np.random.uniform(0, 1), 10)
+    return theta, phi
+
+
+def cirq_sim(q1func, q2func, iterations=1000, sample_count=1):
     all_samples = defaultdict(int)
+    memoized_exp = MemoizedExp(q1func, q2func)
     for i in range(iterations):
-        theta1, phi1 = q1func()
-        theta2, phi2 = q2func()
+        theta1, phi1 = memoized_exp.call(q1func.__name__)
+        theta2, phi2 = memoized_exp.call(q2func.__name__)
         memoized_exp.reset()
 
         q1 = cirq.GridQubit(0, 0)
@@ -97,7 +108,7 @@ def cirq_sim(q1func, q2func, memoized_exp, iterations=1000, sample_count=1):
             rq2 = result.measurements['q2'][rep][0]
             k = "{}{}".format(rq1, rq2)
             all_samples[k] += 1
-    print_all_samples(all_samples)
+    print_all_samples("cirq", all_samples)
 
 
 def math_sim_0(*args, **kwargs):
@@ -109,20 +120,17 @@ def math_sim_1(*args, **kwargs):
 
 
 def cirq_sim_0(*args, **kwargs):
-    memoized_exp = MemoizedExp()
-    cirq_sim(q1func=memoized_exp.gen_exp_for_cirq_0,
-             q2func=memoized_exp.gen_exp_for_cirq_0,
-             memoized_exp=memoized_exp,
+    cirq_sim(q1func=gen_exp_for_cirq_0,
+             q2func=gen_exp_for_cirq_0,
              *args, **kwargs)
 
 
 def cirq_sim_1(*args, **kwargs):
-    memoized_exp = MemoizedExp()
-    cirq_sim(q1func=memoized_exp.gen_exp_for_cirq_0,
+    cirq_sim(q1func=gen_exp_for_cirq_0,
              q2func=gen_exp_for_cirq_1,
-             memoized_exp=memoized_exp,
              *args, **kwargs)
 
 
 if __name__ == "__main__":
+    math_sim_0()
     cirq_sim_0()