diff --git a/__main__.py b/__main__.py
index 16e4e82..7b6bbe4 100644
--- a/__main__.py
+++ b/__main__.py
@@ -1,3 +1,3 @@
-from .week2 import main
+from .week4 import main
 
-main()
\ No newline at end of file
+main()
diff --git a/week3.py b/week3.py
index e5410c3..3d4bda9 100644
--- a/week3.py
+++ b/week3.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import random
 
-from crypto.week1 import BloodType, blood_cell_compatibility_lookup
+from .week1 import BloodType, blood_cell_compatibility_lookup
 
 
 class Protocol:
diff --git a/week4.py b/week4.py
index 48a6eeb..d60c899 100644
--- a/week4.py
+++ b/week4.py
@@ -1,22 +1,13 @@
 # Concept: Create 8 PKs where each represent a bloodtype. Let 7 of them be created by OGen and 1 of them by KeyGen.
 # The one represents our bloodtype. Bob will then encrypt 8 values using these PKs, where each value repredents
 # A truth value, thus either true or false, s.t. each cipher is an entry in the bloodtype comptability matrix.
-
+import math
+import random
 from secrets import SystemRandom
-import numpy as np
-from crypto.week1 import BloodType, blood_cell_compatibility_lookup
 
-# We can't encrypt 0, so we have to index from 1
-convert_bloodtype_to_index = {
-    BloodType.O_NEGATIVE: 1,
-    BloodType.O_POSITIVE: 2,
-    BloodType.A_NEGATIVE: 3,
-    BloodType.A_POSITIVE: 4,
-    BloodType.B_NEGATIVE: 5,
-    BloodType.B_POSITIVE: 6,
-    BloodType.AB_NEGATIVE: 7,
-    BloodType.AB_POSITIVE: 8,
-}
+from .week1 import BloodType, blood_cell_compatibility_lookup
+
+bloodtypes = {b: i for i, b in enumerate(BloodType, start=1)}  # we can't encrypt 0, so we have to index from 1
 
 
 class ElGamal:
@@ -29,7 +20,7 @@ class ElGamal:
 
     def gen_key(self):
         key = SystemRandom().randint(1, self.order)
-        while np.gcd(self.order, key) != 1:
+        while math.gcd(self.order, key) != 1:
             key = SystemRandom().randint(1, self.order)
         return key
 
@@ -67,12 +58,8 @@ class Alice:
 
         self.sk = elgamal.gen_key()
         self.pk = elgamal.gen(self.sk)
-        # There is 100% a better way to get the index, also, it's to avoid encryption 0. Coincidentally, it's
-        # not an issue when we do it, as O- is in the 0th index and this bloodtype can donate to everyone
-        # so the decryption bugging, it not an issue
-        self.b = list(convert_bloodtype_to_index.keys()).index(bloodtype)+1
-        self.fake_pks = [self.elgamal.ogen()
-                         for _ in range(7)]
+        self.b = bloodtypes[bloodtype]
+        self.fake_pks = [self.elgamal.ogen() for _ in range(7)]
 
     def send_pks(self):
         all_pks = self.fake_pks
@@ -80,15 +67,15 @@ class Alice:
         return all_pks
 
     def retrieve(self, ciphers):
-        print(ciphers)
+        #print(ciphers)
         mb = self.elgamal.dec(ciphers[self.b-1])
-        # Bog sends 1 for false, 2 for true, so we have to subtract 1 here
-        return mb-1
+        # Bob sends 1 for false, 2 for true, so we have to subtract 1 here
+        return mb - 1
 
 
 class Bob:
     def __init__(self, bloodtype, elgamal):
-        self.bloodtype = list(convert_bloodtype_to_index.keys()).index(bloodtype)
+        self.bloodtype = bloodtypes[bloodtype]
         self.truth_vals = []
         self.elgamal = elgamal
         self.pks = None
@@ -104,16 +91,66 @@ class Bob:
         ciphers = []
         for idx, truth_val in enumerate(self.truth_vals):
             pk = self.pks[idx]
-            # Bob can't send 0, as it will encrypt to 0
-            c = self.elgamal.enc(truth_val+1, pk)
+            c = self.elgamal.enc(truth_val + 1, pk)  # + 1 since Bob can't send 0, as it will encrypt to 0
             ciphers.append(c)
         return ciphers
 
 
-def run(donor : BloodType, recipient : BloodType):
-    p = 15485863
+def is_prime(n: int, k: int) -> bool:
+    """
+    Miller-Rabin Primality test.
+    Adapted from pseudo-code at https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test.
+
+    :param n: An odd integer to be tested for primality.
+    :param k: The number of rounds of testing to perform.
+    :return: True if n is 'probably prime', False otherwise if n is composite.
+    """
+    # write n as 2r·d + 1 with d odd (by factoring out powers of 2 from n − 1)
+    d = n - 1
+    r = 0
+    while d % 2 == 0:
+        d >>= 1
+        r += 1
+
+    for i in range(k):  # witnessLoop
+        continue_wl = False
+        a = random.randint(2, n - 2)
+        x = pow(a, d, n)
+        if x == 1 or x == n - 1:
+            continue
+        for j in range(r - 1):
+            x = pow(x, 2, n)
+            if x == n - 1:
+                continue_wl = True
+                break
+        if continue_wl:
+            continue
+        return False
+    return True
+
+
+def gen_prime(b: int, k: int = 10) -> int:
+    """
+    Generate strong probable prime by drawing integers at random until one passes the is_prime test.
+    Adapted from pseudo-code at https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test.
+
+    :param b: The number of bits of the result.
+    :param k: The number of rounds of testing to perform.
+    :return: a strong probable prime.
+    """
+    while True:
+        n = random.randint(2**(b-1), (2**b)-1)
+        if n % 2 == 0:
+            continue
+        if is_prime(n, k):
+            return n
+
+
+def run(donor: BloodType, recipient: BloodType):
+    p = gen_prime(128)
     q = 2 * p + 1
     g = SystemRandom().randint(2, q)
+    #print("p:", p, "q:", q, "g:", g)
 
     elgamal = ElGamal(g, q, p)
     alice = Alice(donor, elgamal)
@@ -125,11 +162,7 @@ def run(donor : BloodType, recipient : BloodType):
     return bool(pls)
 
 
-if __name__ == "__main__":
-    #z = run(BloodType.O_POSITIVE, BloodType.A_NEGATIVE)
-    #print(z)
-    #exit()
-
+def main():
     green = 0
     red = 0
     for i, recipient in enumerate(BloodType):