"""
John Rachlin
DS 2000: Intro to Programming with Data


Date: Wed Sep 21 21:22:07 2022

File: collatz_plot.py
    
Description: 
    
    Plot the number of steps it takes before
    the collatz sequence enters the 421421421 cycle
    
    
"""


import matplotlib.pyplot as plt

def main():
    
    plt.figure(figsize=(6,4), dpi=200)
    
    n = 1
    while n < 10000:
        steps = 0
        x = n
        while x != 1:
            
            if x % 2 == 0:
                x = x // 2
            else:
                x = 3 * x + 1
                
            steps = steps + 1
        print(n, steps)
        
        if n % 2 == 0:
            marker_color = 'blue'
        else:
            marker_color = 'red'
            
        plt.plot(n, steps, '.', color=marker_color)
        n = n + 1
        
    plt.grid()
    plt.xlabel('n')
    plt.ylabel('steps')
    plt.title('Collatz steps until 421 cycle. (Red=Odd, Blue=Even)')
    plt.legend()
    plt.savefig('collatz.png')

main()