diff --git a/main.py b/main.py
index 857cd1b..61386cc 100644
--- a/main.py
+++ b/main.py
@@ -1,21 +1,18 @@
 import time
 import numpy as np
 import matplotlib.pyplot as plt
-from tqdm import tqdm
-from typing import Callable, Tuple
+
 from src.cohort_builder import build_cohorts
 from src.cohort_simulator import simulate_cohorts
 from src.param import *
-import concurrent.futures
-
+from src.macro import Macro
 # TODO: @chingyulin: make cohort a class
 
 # The main loop builds up the economy with a large number of cohorts, and simulates the stationary economy forward
 for k in range(Mpaths):
-# def simulate(k, Nc, dt, rho, nu, Vhat, mu_Y, sigma_Y, beta, T_hat):
     s = time.time()
     time_s = time.time()
-    dZt = dt**0.5 * np.random.randn(int(Nt - 1))
+    macro = Macro(dt, T_cohort, mu_Y, sigma_Y)
     (
         IntVec,
         Xt,
@@ -25,7 +22,7 @@
         tau,
         MaxThetaDelta_s_t,
         invest_tracker_keep,
-    ) = build_cohorts(dZt, Nc, dt, rho, nu, Vhat, mu_Y, sigma_Y, beta, T_hat, mode1)
+    ) = build_cohorts(dt, rho, nu, Vhat, beta, T_hat, mode1)
 
     (
         IntVec_drop,
@@ -37,8 +34,6 @@
         MaxThetaDelta_s_t_drop,
         invest_tracker,
     ) = build_cohorts(dZt, Nc, dt, rho, nu, Vhat, mu_Y, sigma_Y, beta, T_hat, mode2)
-    # if time.time() - time_s > time_tolerance:
-    #     print(f"It takes more than {time_tolerance}s to build up the cohorts")
 
     dZforbias = np.diff(Zt)  # dZt used in the build_cohorts function
     biasvec = dZforbias[-Npre:]
diff --git a/src/cohort_builder.py b/src/cohort_builder.py
index d68836d..7e7d10d 100644
--- a/src/cohort_builder.py
+++ b/src/cohort_builder.py
@@ -1,73 +1,49 @@
 import numpy as np
 from src.solver import bisection, solve_theta
 from tqdm import tqdm
-from typing import Tuple
 from src.stats import post_var
-
+from src.macro import Macro
 
 def build_cohorts(
-    dZt: np.ndarray,
-    Nc: int,
+    macro: Macro,
     dt: float,
+    Nc: int,
     rho: float,
     nu: float,
     Vhat: float,
-    mu_Y: float,
-    sigma_Y: float,
     beta: float,
     T_hat: float,
     mode: str
-) -> Tuple[
-    # np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    # np.ndarray,
-    np.ndarray,
-    np.ndarray,
-]:
+):
     """builds up a sufficiently large set of cohorts in the economy, view each cohort as one agent with a constantly shrinking size
 
     Args:
-        dZt (np.ndarray): random shocks of aggregate output for each period, shape (Nc-1, )
         Nc (int): number of periods  = number of cohorts in the economy
         dt (float): unit of time
         rho (float): rho, discount factor
         nu (float): birth / death rate, each cohort starts at size nu and shrinks at speed of nu
         Vhat (float): initial variance of beliefs
-        mu_Y (float): mean of aggregate output growth
-        sigma_Y (float): sd of aggregate output growth
         beta (float): initial consumption of the newborn agents
         T_hat (float): pre-trading years
         mode (str): describes the mode
 
     Returns:
-        # DeltaConditional (np.ndarray): consumption weighted aggregate max(delta_s_t, -theta_t), as in eq(19), shape(Nc, )
         IntVec (np.ndarray): ~similar to consumption share, shape(Nc, )
         Xt (np.ndarray): xi_t * Yt, shape(Nc, )
         Delta_s_t (np.ndarray): bias, shape(Nc, )
         Yt (np.ndarray): aggregate output, shape(Nc, )
-        Zt (np.ndarray): cumulated shocks, shape(Nc, )
         # consumptionshare (np.ndarray): shape(Nc, )
         tau (np.ndarray): t-s, shape(Nc, )
         MaxThetaDelta_s_t (np.ndarray): max(delta_s_t, -theta_t), shape(Nc, )
-        #TODO: @chingyulin: use NamedTuple for the return
     """
-
     Npre: int = int(T_hat / dt)  # Number of pre-trading observations
-    Zt = np.insert(np.cumsum(dZt), 0, 0)  # cumulated shocks, Nc * 1
-    yg = (mu_Y - 0.5 * sigma_Y**2) * dt + sigma_Y * dZt  # output in log, (Nc - 1) *1, eq(1)
-    Yt = np.insert(np.exp(np.cumsum(yg)), 0, 1)  # output, Nc *1
-    # DeltaConditional = np.zeros(Nc)
     Delta_s_t = np.zeros(1)  # belief bias, eq(3)
     MaxThetaDelta_s_t = np.zeros(1)  # disagreement, eq(11)
     Xt = np.ones(Nc) * nu * beta  # similar to consumption share, similar to eq(18)
     IntVec = nu * beta  # consumption share of a newborn cohort
     # TODO: @chingyulin: tau can allocate the memory
     tau = np.zeros(1)  # t-s
-    tau[0] = dt
+    tau[0] = macro.dt
     reduction = np.exp(-nu * dt)  # cohort size shrink at this rate
     theta = np.zeros(Nc)  # market price of risk
     invest_tracker = np.ones(Npre)
@@ -75,7 +51,7 @@ def build_cohorts(
     #for i in tqdm(range(1, Npre)):
         Part = IntVec * np.exp(
             -(rho + 0.5 * MaxThetaDelta_s_t * MaxThetaDelta_s_t) * dt
-            + MaxThetaDelta_s_t * dZt[i - 1]
+            + MaxThetaDelta_s_t * macro.dZt[i - 1]
         )  # Consumption of each cohort, eq(16), where eta_s_t / eta_s_s follows eq(11)
         if i == 1:  # only one cohort in the economy
             Xt[i] = Part
@@ -93,15 +69,15 @@ def build_cohorts(
         consumptionshare = IntVec / Xt[i]  # consumption share
 
         # update beliefs
-        dDelta_s_t = (post_var(sigma_Y, Vhat, tau) / sigma_Y**2) * (
-            -Delta_s_t * dt + dZt[i - 1]
+        dDelta_s_t = (post_var(macro.sigma_Y, Vhat, tau) / macro.sigma_Y**2) * (
+            -Delta_s_t * dt + macro.dZt[i - 1]
         )  # from eq(5)
         if i < Npre:
             # TODO: @chingyulin: this can be optimized
             Delta_s_t = Delta_s_t + dDelta_s_t
             Delta_s_t = np.append(Delta_s_t, 0)  # newborns begin with 0
         else:
-            DELbias = np.sum(dZt[int(i - Npre) : i]) / T_hat
+            DELbias = np.sum(macro.dZt[int(i - Npre) : i]) / T_hat
 
             Delta_s_t += dDelta_s_t
             # TODO: @chingyulin: this can be optimized
@@ -143,151 +119,12 @@ def build_cohorts(
                 )  # update max(Delta_s_t, -theta)
 
     return (
-        # DeltaConditional,
         IntVec,
         Xt,
         Delta_s_t,
-        Yt,
-        Zt,
-        # consumptionshare,
+        macro.Yt,
+        macro.Zt,
         tau,
         MaxThetaDelta_s_t,
         invest_tracker,
-    )
-
-
-#####################################################################################
-
-def build_cohorts_complete_market(
-    dZt: np.ndarray,
-    Nc: int,
-    dt: float,
-    rho: float,
-    nu: float,
-    Vhat: float,
-    mu_Y: float,
-    sigma_Y: float,
-    beta: float,
-    T_hat: float,
-) -> Tuple[
-    np.ndarray,
-    np.ndarray,
-    # np.ndarray,
-    # np.ndarray,
-    # np.ndarray,
-    # np.ndarray,
-    # np.ndarray,
-    # np.ndarray,
-]:
-    """builds up a sufficiently large set of cohorts in the economy, view each cohort as one agent with a constantly shrinking size
-    complete market version
-    run this function along with the incomplete version above,
-    as this function returns results that are used in comparison with the previous results
-
-    Args:
-        dZt (np.ndarray): random shocks of aggregate output for each period, shape (Nc-1, )
-        Nc (int): number of periods  = number of cohorts in the economy
-        dt (float): unit of time
-        rho (float): rho, discount factor
-        nu (float): birth / death rate, each cohort starts at size nu and shrinks at speed of nu
-        Vhat (float): initial variance of beliefs
-        mu_Y (float): mean of aggregate output growth
-        sigma_Y (float): sd of aggregate output growth
-        beta (float): initial consumption of the newborn agents
-        T_hat (float): pre-trading years
-
-    Returns:
-        # DeltaConditional (np.ndarray): consumption weighted aggregate max(delta_s_t, -theta_t), as in eq(19), shape(Nc, )
-        IntVec (np.ndarray): ~similar to consumption share, shape(Nc, )
-        Xt (np.ndarray): xi_t * Yt, shape(Nc, )
-        # Delta_s_t (np.ndarray): bias, shape(Nc, )
-        # Yt (np.ndarray): aggregate output, shape(Nc, )
-        # Zt (np.ndarray): cumulated shocks, shape(Nc, )
-        # consumptionshare (np.ndarray): shape(Nc, )
-        # tau (np.ndarray): t-s, shape(Nc, )
-        # MaxThetaDelta_s_t (np.ndarray): max(delta_s_t, -theta_t), shape(Nc, )
-    """
-
-    Npre: int = int(T_hat / dt)  # Number of pre-trading observations
-
-    # Zt = np.insert(np.cumsum(dZt), 0, 0)  # cumulated shocks, Nc * 1
-    # yg = (mu_Y - 0.5 * sigma_Y**2) * dt + sigma_Y * dZt  # output in log, (Nc - 1) *1, eq(1)
-    # Yt = np.insert(np.exp(np.cumsum(yg)), 0, 1)  # output, Nc *1
-    # DeltaConditional = np.zeros(Nc)
-    Delta_s_t = np.zeros(1)  # belief bias, eq(3)
-    # MaxThetaDelta_s_t = np.zeros(1)  # disagreement, eq(11)
-    Xt = np.ones(Nc) * nu * beta  # similar to consumption share, similar to eq(18)
-    IntVec = nu * beta  # consumption share of a newborn cohort
-
-    tau = np.zeros(1)  # t-s
-    tau[0] = dt
-    reduction = np.exp(-nu * dt)  # cohort size shrink at this rate
-    theta_t = np.zeros(Nc)  # market price of risk
-    for i in tqdm(range(1, Nc)):
-        Part = IntVec * np.exp(
-            -(rho + 0.5 * Delta_s_t * Delta_s_t) * dt
-            + Delta_s_t * dZt[i - 1]
-        )  # Consumption of each cohort, eq(16), where eta_s_t / eta_s_s follows eq(11)
-        if i == 1:  # only one cohort in the economy
-            Xt[i] = Part
-            # DeltaConditional[i] = Part * MaxThetaDelta_s_t
-        else:  # more cohorts
-            Xt[i] = np.sum(Part)  # total consumption
-            #DeltaConditional[i] = (
-                #np.sum(Part * MaxThetaDelta_s_t) / Xt[i]
-            #)  # eq(19), consumption weighted max(Delta_s_t, -theta)
-
-        IntVec = reduction * Part
-        IntVec = np.append(
-            IntVec, beta * (1 - reduction) * Xt[i]
-        )  # updated consumption, add a newborn cohort
-        consumptionshare = IntVec / Xt[i]  # consumption share
-
-        # update beliefs
-        dDelta_s_t = (post_var(sigma_Y, Vhat, tau) / sigma_Y**2) * (
-            -Delta_s_t * dt + dZt[i - 1]
-        )  # from eq(5)
-        if i < Npre:
-            # TODO: @chingyulin: this can be optimized
-            Delta_s_t = Delta_s_t + dDelta_s_t
-            Delta_s_t = np.append(Delta_s_t, 0)  # newborns begin with 0
-        else:
-            DELbias = np.sum(dZt[int(i - Npre) : i]) / T_hat
-
-            Delta_s_t += dDelta_s_t
-            # TODO: @chingyulin: this can be optimized
-            Delta_s_t = np.append(
-                Delta_s_t, DELbias
-            )  # newborns begin with available earlier observations
-
-        # update tau
-        tau += dt
-        tau = np.append(tau, 0)  # TODO: @chingyulin: this can be optimized
-
-        # find the market clearing theta, given beliefs and consumption shares
-        # need a large enough number of cohorts to make the distribution of beliefs reasonably continuous
-        # if i < Npre:
-        #     MaxThetaDelta_s_t = (
-        #         Delta_s_t  # relax the short-sale constraint in the beginning
-        #     )
-        # else:
-        #     lowest_bound = -np.max(Delta_s_t)  # absolute lower bound for theta
-        #     # Should it's put as a configurable parameter?
-        #     theta_t[i] = bisection(
-        #         solve_theta, lowest_bound, 10, consumptionshare, Delta_s_t, sigma_Y
-        #     )  # solve for theta
-        #     MaxThetaDelta_s_t = np.maximum(
-        #         -theta_t[i], Delta_s_t
-        #     )  # update max(Delta_s_t, -theta)
-
-    return (
-        # DeltaConditional,
-        IntVec,
-        Xt,
-        # Delta_s_t,
-        # Yt,
-        # Zt,
-        # consumptionshare,
-        # tau,
-        # MaxThetaDelta_s_t,
-    )
+    )
\ No newline at end of file
diff --git a/src/cohorts.py b/src/cohorts.py
new file mode 100644
index 0000000..db2ac94
--- /dev/null
+++ b/src/cohorts.py
@@ -0,0 +1,25 @@
+import numpy as np
+
+class Believe:
+    def __init__(self, Vhat) -> None:
+        self.Delta_s_t = np.zeros(1)
+        self.Vhat = Vhat
+
+class BaseCohorts:
+
+    def __init__(self, Nc: int, nu: float, dt: float, beta: float, Vhat: float):
+        self.Nc = Nc
+        self.reduction = np.exp(-nu * dt)
+        self.IntVec = nu * beta
+        self.believe = Believe(Vhat)
+
+    def evolve(self): 
+        """Evolve one time unit"""
+        self.update_believe()
+        self.decide_investment()
+
+    def update_believe(self):
+        ...
+
+    def decide_investment(self): 
+        ...
diff --git a/src/macro.py b/src/macro.py
new file mode 100644
index 0000000..fc45d83
--- /dev/null
+++ b/src/macro.py
@@ -0,0 +1,87 @@
+import numpy as np
+
+class Macro:
+    def __init__(self, dt: float, T_cohort: float, mu_Y: float, sigma_Y: float) -> None:
+        """Macro-economic
+
+        Args:
+            dt (float): time incremental
+            T_cohort (float): time horizon to keep track of cohorts
+            mu_Y (float): mean of aggregate output growth
+            sigma_Y (float): sd of aggregate output growth
+        """
+        self.dt = dt
+        self.T_cohort = T_cohort
+        self.Nt = int(T_cohort / dt)
+        self.dZt =  self.dt**0.5 * np.random.randn(self.Nt - 1)
+        self.Zt = np.insert(np.cumsum(self.dZt), 0, 0)
+        self.mu_Y = mu_Y
+        self.sigma_Y = sigma_Y
+        self.yg = (mu_Y - 0.5 * sigma_Y**2) * dt + sigma_Y * self.dZt
+        self.Yt = np.insert(np.exp(np.cumsum(self.yg)), 0, 1)
+
+    
+
+
+
+# class BaseCohorts:
+#     a = "C"
+
+#     def __init__(self, t: float, T: float) -> None:
+#         self.t  = t
+#         self.T = T
+#         self.believe = ()
+#         self.ratio = ()
+
+#     @staticmethod
+#     def print():
+#         print("HH")
+
+#     @classmethod
+#     def print(cls):
+#         cls.a
+
+
+#     def one_year(self):
+#         10 ** self.update_believe()
+#         self.invest()
+#         self.t += 1
+
+#         return self.t
+
+#     def update_believe(self) -> int:
+#         self.believe
+#         return self.believe
+
+#     def invest(self):
+#         print("Base")
+#         self.ratio
+
+# class NewBelieveCohort(BaseCohorts):
+#     def update_believe(self) -> int:
+#         print("Do something different.")
+#         return self.believe + 0.5
+
+# class NewerBelieveCohort(NewBelieveCohort):
+#     ...
+
+# class DropCohort(BaseCohorts):
+
+#     def __init__(self, t: float, T: float, b: float) -> None:
+#         super().__init__(t, T)
+#         self.b = b
+
+#     def invest(self):
+#         print("Drop")
+
+#         self.ratio
+
+# if __name__ == "__main__":
+
+#     macro = Macro(z_t = 1, sigma_y=1)
+#     print(macro.y_t)
+
+#     cohort = DropCohort(0, 10, 1)
+
+#     for i in range(100):
+#         t = cohort.one_year()
\ No newline at end of file