import numpy as np
import scipy as scp

from matplotlib import pyplot as plt

from RubberBand import RubberBand

USE_SAVED = True

NUM_BANDS = int(1E6)
N = 100

if USE_SAVED:
    lengths = np.load("data/lengths.npy")
else:
    lengths = np.zeros((NUM_BANDS,))
    
    for i in range(NUM_BANDS):
        band = RubberBand(N, a=1)
        lengths[i] = band.length/band.a

# histogram
values, bin_edges = np.histogram(lengths, bins=int(2*N), range=(-N, N))
p_mc = values / values.sum() # normalise histogram

bin_centers = bin_edges[:-1] + np.diff(bin_edges)/2

def P(N, L_over_a):
    """
    P(L/a) = Omega(N, n^+) / 2^N
    where
    n^+ = (L/a + N)/2 (see eq. 1)
    """
    return scp.special.binom(N, (N+L_over_a)/2) / 2**N

p_true = P(N, bin_centers-.5)

chi2 = np.sum((p_mc - p_true)**2 / p_true)

# plot hist vs true dist.
plt.plot(bin_centers-.5, p_true, label="$P(L)$")
plt.step(bin_centers, p_mc, color='r', alpha = 0.5, label="$\hat{P}(L)$")
plt.text(-45, 0.07, f"$\chi^2 = {round(chi2, 3)}$\n({len(lengths):.0E} samples)")
plt.legend()
plt.xlabel("$L/a$")
plt.ylabel("P($L/a$)")
xlim = (-50, 50)
plt.xlim(xlim)

# ratios
p_ratio = p_mc / p_true
plt.figure()
plt.scatter(bin_centers, p_ratio, marker='x')
plt.hlines(1, *xlim, color='r', linestyle='--')
plt.xlim(xlim)
plt.ylim(0.5, 1.5)
plt.xlabel("$L/a$")
plt.ylabel("$\hat{P}(L)/P(L)$")
plt.show()

np.save("data/lengths.npy", lengths)