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
bin_edges = np.arange(-N-1, N+1, 2)
values, _ = np.histogram(lengths, bins=bin_edges)
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

# Calculate p_true, scale with NUM_BANDS,
# mask low statistics bins,
# calculate chi^2/ndf

p_true = P(N, bin_centers)
p_true_scaled = p_true*np.sum(values)
values_masked = values[values > 5]
p_true_scaled_masked = p_true_scaled[values > 5]
chi2 = np.sum((values_masked - p_true_scaled_masked)**2 / p_true_scaled_masked)
chi2_over_ndf =  chi2 / (len(values_masked)-1)

# plot hist vs true dist.
plt.plot(bin_centers, p_true, label="$P(L)$")
plt.bar(bin_centers, p_mc, color='r', alpha = 0.5, label="$\hat{P}(L)$")
plt.text(-45, 0.07, f"$\chi^2/ndf = {round(chi2_over_ndf, 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.xlabel("$L/a$")
plt.ylabel("$\hat{P}(L)/P(L)$")
plt.show()

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