import numpy as np
from matplotlib import pyplot as plt

from RubberBand import RubberBand

def generate_ensemble(f):
    """
    generate an ensemble of RubberBands with a force f applied
    """
    
    lengths = np.zeros((NUM_BANDS,))
    
    for i in range(NUM_BANDS):
        band = RubberBand(N, a=a, kBT=kBT, force=f)
        lengths[i] = band.length/band.a
    
    return lengths

def calculate_k_eff_theory(N, a, kBT):
    return N*a**2/kBT
    
def true_expectation_values(f, N, a, kBT):
    """
    compute 'true' expectation value <L> for both full analytical one and
    using the small force limit for the linear region
    """
    L_analytical = N*a*np.tanh(a*f/kBT)
    L_analytical_linear = f*calculate_k_eff_theory(N, a, kBT)
    
    return L_analytical, L_analytical_linear

def plot_true_L_vs_sim_L(forces, means, stdevs, params):
    L_true, L_true_linear = true_expectation_values(forces, *params)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(5, 9))
    
    # Plot 1: true L vs approx L in entire force range
    ax1.plot(forces, L_true, color='red',
             label="$\langle L \\rangle$ (analytical)")
    # ax1.plot(forces, L_true_linear, color='red', alpha=1,
             # label="$\langle L \\rangle$ (analytical, small force approximation)")
    ax1.plot(forces, means-stdevs, color='blue', alpha=.5)
    ax1.plot(forces, means+stdevs, color='blue', alpha=.5)
    ax1.fill_between(forces, means-stdevs, means+stdevs, color='blue', alpha=0.1)
    ax1.plot(forces, means, color='blue', linestyle=':', linewidth=2,
             label=f"$\langle \hat{{L}} \\rangle$ (M={NUM_BANDS:.0E})")
    ax1.set_xlabel("$f$ [arb.]")
    ax1.set_ylabel("$\langle L \\rangle (f)$")
    ax1.legend()
    
    # Plot 2: true L vs approx L in small force range
    ax2.plot(forces, L_true_linear, color='red',
             label="$\langle L \\rangle$ (small force approximation)")
    ax2.plot(forces, means-stdevs, color='blue', alpha=.5)
    ax2.plot(forces, means+stdevs, color='blue', alpha=.5)
    ax2.fill_between(forces, means-stdevs, means+stdevs, color='blue', alpha=0.1)
    ax2.plot(forces, means, color='blue', linestyle=':', linewidth=2,
             label=f"$\langle \hat{{L}} \\rangle$ (M={NUM_BANDS:.0E})")
    ax2.set_xlim(0, 0.5)
    ax2.set_ylim(0, 75)
    ax2.set_xlabel("$f$ [arb.]")
    ax2.set_ylabel("$\langle L \\rangle (f)$")
    ax2.legend()
    plt.tight_layout()
    
    # plot 3: abs difference between true and approx L
    plt.figure()
    plt.plot(forces, np.abs(L_true-means), 
             label="$\langle L \\rangle$ (analytical)")
    plt.plot(forces, np.abs(L_true_linear-means),
             label="$\langle L \\rangle$ (small force approximation)")
    plt.yscale('log')
    plt.xlabel("$f$ [arb.]")
    plt.ylabel("$| \langle L \\rangle - \langle \hat{L} \\rangle |$")
    plt.legend()
    plt.tight_layout()
    
    plt.show()
    
def estimate_k_eff(forces, means, upper_lim=0.5):
    x = forces[forces < upper_lim]
    y = means[forces < upper_lim]
    m, b = np.polyfit(x, y, 1)
    return m

if __name__ == "__main__":
    # simulation parameters
    USE_SAVED = True
    NUM_BANDS = int(1E3)
    N = 100
    a = 1.
    kBT = 1.
    
    params = (N, a, kBT)
    
    forces = np.arange(0., 5, 0.01)
    
    if USE_SAVED:
        means = np.load("data/task3_means.npy")
        stdevs = np.load("data/task3_stdevs.npy")
    else:
        means = np.zeros(len(forces))
        stdevs = np.zeros(len(forces))
        
        for i, f in enumerate(forces):
            L_i = generate_ensemble(f)
            means[i] = np.mean(L_i)
            stdevs[i] = np.std(L_i)
            
        np.save("data/task3_means.npy", means)
        np.save("data/task3_stdevs.npy", stdevs)
    
    # generate the ensembles and compute values
    plot_true_L_vs_sim_L(forces, means, stdevs, params)
    
    lim = 0.3
    k_eff_est = estimate_k_eff(forces, means, upper_lim=lim)
    print(f"""Analytical k_eff v.s. estimate from fit:
          analytical (N * a^2/kBT):     k_eff = {calculate_k_eff_theory(*params)}
          estimate (f < {lim}):         k_eff = {round(k_eff_est, 2)}
          """)