from math import factorial
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from matplotlib import cm
import time

def oppg1():
    n_list = [int(n) for n in np.logspace(0,5,50)]
    t_list = np.zeros(len(n_list))

    n_fac = 0 #Tar med denne for at den skal overskrive variabelen også på første iterasjon
    t = 0 #Samme for t
    for i, n in enumerate(n_list):
        t0 = time.process_time()
        n_fac = factorial(n)
        t = time.process_time() - t0
        t_list[i] = t

    plt.plot(n_list, t_list)
    plt.xlabel('n')
    plt.ylabel('t [s]')
    plt.title(r'Time to calculate $n!$')
    plt.yscale('log')
    plt.xscale('log')
    plt.savefig('oppg1')
    plt.show() #Åpenbart ikke mulig å beregne n! for makroskopiske systemer med 1e24 + antall partikler

def oppg3():
    cmap = cm.get_cmap('viridis')
    stirling = lambda i: i * np.log(i) - i
    lnW = lambda N, n : stirling(N) - stirling(n) - stirling(N - n) #W = N! / ( n!(N - n)! )

    n, N = sp.symbols('n, N')
    lnW_symbolic = - (n * sp.log(n) - n) - ((N - n) * sp.log(N - n) - (N - n))
    dlnW_dn = sp.diff(lnW_symbolic, n)
    n_max = sp.solve(dlnW_dn, n)[0] #finner max  multiplisitet som funksjon av antall gitterpunkter


    N_list = [int(i) for i in np.logspace(1, 6, 10)]

    for N_val in N_list:
        color = cmap(np.log10(N_val)/np.log10(max(N_list)))
        n_list = np.array([int(n) for n in np.logspace(1, np.log10(N_val), 100)])
        n_max_val = int(n_max.subs({N : N_val}))
        S_max = lnW(N_val, n_max_val)
        S = lnW(N_val, n_list)
        plt.plot(n_list, S, label=round(np.log10(N_val), 0), color=color)
        plt.scatter(n_max_val, S_max, color=color)

    plt.legend(title = r'$\log{N}$')
    plt.ylabel(r'$\frac{S}{k_B}$', fontsize = 14)
    plt.xlabel('n particles')
    plt.xscale('log')
    plt.yscale('log')
    plt.title('Multiplicity as a function of lattice sites and number of particles')
    plt.savefig('oppg3')
    plt.show()

oppg1()
oppg3()