import numpy as np
import matplotlib.pyplot as plt

### Modular arithmetic and powers of numbers

def pow(n,p):
    """Gives all the powers of number n modulo p.
    Inputs:
    - n: Int - Number whose powers are computed
    - p: Int - Number modulo with these powers.
    Gives list of powers of n modulo p and length of the list."""
    k = n%p
    L = [1]
    while k not in L:
        L += [k]
        k = (n*k)%p
    return L , len(L)

def graph_pow(N , save=False):
    """Gives graph of pow(n,p) for 1 <= n,p <= N.
    Inputs:
    - N: Int - Max number selected.
    - save: Boolean - Saves the figure or not. Default: False."""

    M = np.zeros([N,N])

    for i in range(N):
        for j in range(N):
            M[i,j] = pow(i+1,j+1)[1]

    plt.figure()
    #plt.imshow(np.log(M) , cmap="jet" , extent=(1,N+1,N+1,1))
    plt.imshow(M , cmap="jet" , extent=(1,N+1,N+1,1))
    plt.colorbar()
    plt.xlabel("$p$")
    plt.ylabel("$n$")
    plt.title("Numbers of powers of n modulo p")
    if save == True:
        plt.savefig("power_numbers_N="+str(N)+".pdf")
    if save == False:
        plt.show()
    pass
