import numpy as np
import scipy as scp

from matplotlib import pyplot as plt

from RubberBand import RubberBand

# load
lengths = np.load("data/lengths.npy")
NUM_BANDS = len(lengths)
N = 100

dummy_band = RubberBand(N=N) # create a dummy band to access w

def calculate_mu_eff(w):
    w_tilde = w / np.sum(w)
    return 1/np.sum(w_tilde**2)

def P_weighted(L_over_a, f, dummy_band):
    """
    Weighted true probablity density function.
    """
    N = dummy_band.N
    weights = dummy_band.w(f, L_over_a)
    z = scp.special.binom(N, (N+L_over_a)/2) * weights
    return z / np.sum(z)

fig_dist, axes_dist = plt.subplots(3, 2, figsize=(12, 10))
fig_ratio, axes_ratio = plt.subplots(3, 2, figsize=(12, 10))

axes_dist = axes_dist.flatten()
axes_ratio = axes_ratio.flatten()

mu_effs = []
forces = [+0.01, +0.05, +0.1, +0.25, +0.5, +1.0]

for i, f in enumerate(forces):

    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
    
    # weighted histogram
    weights = dummy_band.w(f, bin_centers)
    values_w = values * weights
    p_mc_w = values_w / values_w.sum()
    
    p_true_w = P_weighted(bin_centers, f, dummy_band)
    
    # compute effective samples due to reweighting
    mu_eff = calculate_mu_eff(weights)
    mu_effs.append(mu_eff)
    
    # plot stuff
    ax_dist = axes_dist[i]
    ax_dist.plot(bin_centers, p_true_w, label="$P(L)$ (with $f$)")
    ax_dist.bar(bin_centers, p_mc, color='g', linestyle='--', alpha = 0.2, label="$\hat{P}(L)$ (no $f$)")
    ax_dist.bar(bin_centers, p_mc_w, color='r', alpha = 0.5, label="$\hat{{P}}(L)$ (with $f$)")
    if i == 0:
        ax_dist.legend()
      
    ax_dist.set_xlabel("$L/a$")
    ax_dist.set_ylabel("P($L/a$)")
    
    ax_dist.text(0.03,0.83, f"$f={f}$\n$\mu_{{eff}}={round(mu_eff, 1)}$ %",
                 transform=ax_dist.transAxes)
    
    xlim = (-50, 50)

    # ratios
    p_ratio = p_mc_w / p_true_w
    
    ax_ratio = axes_ratio[i]
    ax_ratio.text(0.03,0.83, f"$f={f}$\n$\mu_{{eff}}={round(mu_eff, 1)}$ %",
                 transform=ax_ratio.transAxes)
    
    ax_ratio.scatter(bin_centers, p_ratio, marker='x')
    ax_ratio.hlines(1, *xlim, color='r', linestyle='--')
    ax_ratio.set_xlim(xlim)
    ax_ratio.set_xlabel("$L/a$")
    ax_ratio.set_ylabel("$\hat{P}(L)/P(L)$")

plt.tight_layout()

plt.figure()
plt.plot(forces, mu_effs)
plt.xlabel("$f$ [arb.]")
plt.ylabel("$\mu_{eff}$ [%]")
plt.show()