diff --git a/src/main.py b/src/main.py
index 3ed909f..d7640f3 100644
--- a/src/main.py
+++ b/src/main.py
@@ -30,8 +30,30 @@ def parse_arguments():
 
     # TODO pokracovani
     
-    block_transactions = [COINBASE_TRANSACTION] + mempool.valid_transactions
-    
+
+    #block_transactions = [COINBASE_TRANSACTION] + mempool.valid_transactions
+
+    # Initialize an empty list for transactions
+    block_transactions = [COINBASE_TRANSACTION]
+
+    # Initialize total weight and total fees
+    total_weight = 0
+    total_fees = 0
+
+    # Set the maximum block weight
+    max_block_weight = 4000000
+
+    # Sort the transactions by the fee in descending order
+    transactions_sorted_by_fee = sorted(mempool.valid_transactions, key=lambda tx: tx.fee, reverse=True)
+
+    for tx in transactions_sorted_by_fee:
+        tx_weight = tx.calculate_weight()
+        if total_weight + tx_weight > max_block_weight:
+            break
+        block_transactions.append(tx.json_transaction)
+        total_weight = total_weight + tx_weight
+        total_fees = total_fees + tx.fee
+
     transaction_hashes = [calculate_txid(COINBASE_TRANSACTION)] + [calculate_txid(json_transaction) for json_transaction in block_transactions[1:]]
     block_hash = block_mining(transaction_hashes).hex()
 
diff --git a/src/mempool.py b/src/mempool.py
index 1231f62..fa11785 100644
--- a/src/mempool.py
+++ b/src/mempool.py
@@ -7,4 +7,4 @@ def __init__(self, root_dir):
         self.root_dir = root_dir
         self.transaction_files = [os.path.join(self.root_dir, file) for file in os.listdir(self.root_dir) if file.endswith('.json')]
         self.transactions = [Transaction(file) for file in self.transaction_files]
-        self.valid_transactions = [transaction.json_transaction for transaction in self.transactions if transaction.is_valid()]
+        self.valid_transactions = [transaction for transaction in self.transactions if transaction.is_valid()]
diff --git a/src/serialize.py b/src/serialize.py
index db84e33..d3bac3c 100644
--- a/src/serialize.py
+++ b/src/serialize.py
@@ -82,6 +82,9 @@ def serialize_transaction(transaction, index=-1, sighash_type=1, segwit=False):
     # witness
     if segwit:
         for tx_in in inputs:
+            if "witness" not in tx_in:
+                break
+
             out += [encode_varint(len(tx_in["witness"]))]
 
             for item in tx_in["witness"]:
diff --git a/src/transaction.py b/src/transaction.py
index 4f24bbf..375fa2e 100644
--- a/src/transaction.py
+++ b/src/transaction.py
@@ -39,6 +39,7 @@ def __init__(self, transaction_json_file):
             self.vin = json_transaction['vin']
             self.vout = json_transaction['vout']
             self.json_transaction = json_transaction
+            self.fee = 0
         else:
             # TODO jestli nejakej error
             print('Invalid transaction syntax')
@@ -90,12 +91,20 @@ def check_input_output_sum(self):
         for output in self.vout:
             output_sum = output_sum + output['value']
         
+        self.fee = input_sum - output_sum
+
         # Output sum can't be greater than the input sum.
         if input_sum < output_sum:
             return False
         
         return True
 
+    def calculate_weight(self):
+        base_size = len(serialize_transaction(self.json_transaction))
+        total_size = len(serialize_transaction(self.json_transaction, segwit=True))
+
+        return base_size * 1.8 + total_size
+
     def valid_input(self, vin_idx, vin):
         if vin.get("is_coinbase", False):
             return False