import numpy as np  
from collections import defaultdict  
import matplotlib.pyplot as plt

def simulate(alpha, num_blocks=100000):
    honest_weight = 0  # cumulative sum of weight of honest chain
    attacker_weight = 0  # cumulative sum of weight of attackers chain
    attackers_blocks = 0
    honest_blocks = 0
    attacker_reward = 0
    honest_reward = 0

    for _ in range(num_blocks):
                
        r = np.random.rand()
        weight = np.log(r)
        
        if r < alpha:  # Attacker mines
            attacker_weight += weight
            attackers_blocks += 1
            
            if attacker_weight < honest_weight:
            
                # attacker reveals branch when it's going to be selected as the best branch
                attacker_reward += attackers_blocks
                
                # honest miners lose their rewards
                honest_reward -= honest_blocks
                
                # start again
                attackers_blocks = honest_blocks = attacker_weight = honest_weight = 0
                
            
        else:  # Honest miner mines
            honest_weight += weight
            honest_reward += 1
            honest_blocks += 1
    
    print(alpha, attacker_reward, honest_reward)
    if honest_reward > 0:
        return attacker_reward / honest_reward
    else:
        return num_blocks

# selfish mining as explained in "Majority is not Enough: Bitcoin Mining is Vulnerable" https://arxiv.org/pdf/1311.0243
# note that gamma is optimistic so this is a lower bound
def simulateBitcoin(alpha, num_blocks=100000, gamma=0.5):
    honest_weight = 0  # Length of honest chain
    attacker_weight = 0  # Length of attacker's private fork
    attacker_reward = 0
    honest_reward = 0
    deltaPrev = 0

    for _ in range(num_blocks):
        deltaPrev = attacker_weight - honest_weight
        
        if np.random.rand() < alpha:  # Attacker mines
            attacker_weight += 1
            
            if deltaPrev == 0 and attacker_weight == 2:
                attacker_reward += attacker_weight
                attacker_weight=honest_weight=0
            
        else:  # Honest miner mines
            honest_weight += 1
            
            if deltaPrev<=0:
                honest_reward += honest_weight
                attacker_weight = honest_weight = 0
            elif deltaPrev==1:
                # miners chose at random which fork to mine on based on gamma parameter
                g = np.random.rand()
                if g <= gamma:
                    attacker_reward += 1
                else:
                    honest_reward += 1
                
                attacker_weight=honest_weight=0
            elif deltaPrev>=2:
                attacker_reward += deltaPrev+1
                attacker_weight = honest_weight = 0       
                
        
    print(alpha, attacker_reward, honest_reward)
    return attacker_reward / honest_reward
    
alphas = np.linspace(0.05, 0.5, 40)  # 5% to 50% attacker hashing power
results_poem = [simulate(alpha) for alpha in alphas]  
results_bitcoin = [simulateBitcoin(alpha) for alpha in alphas]  

# Plot results  
plt.plot(alphas, results_poem, label="PoEM")  
plt.plot(alphas, results_bitcoin, label="Bitcoin")  
plt.axhline(y=1, color="red", linestyle="--", label="Attack break-even cost")  
plt.xlabel("Attacker Hashrate (α)")  
plt.ylabel("Attacker Reward / Honest Reward")  
plt.ylim([0, 1.2])
plt.legend()  
plt.show()