fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plot_mu_trajectory(ax, mu_modes[25600][0][:2*10**6])
ax.set_xlabel('generation', fontsize=28);
ax.set_ylabel('mode of the mutation rate, $\mu_{mode}$', fontsize=28);
plt.savefig('mu_mode_trajectoryK25600.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plot_mu_trajectory(ax, mu_modes[400][0][:2*10**4])
ax.set_xlabel('generation', fontsize=28);
ax.set_ylabel('mode of the mutation rate, $\mu_{mode}$', fontsize=28);
plt.savefig('mu_mode_trajectoryK400.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plot_mu_trajectory(ax, mu_modes[1600][0][:2*10**5])
ax.set_xlabel('generation', fontsize=28);
ax.set_ylabel('mode of the mutation rate, $\mu_{mode}$', fontsize=28);
plt.savefig('mu_mode_trajectoryK1600.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plt.plot(f_maxes[200][1][::100000])

"""
Explanation: Plots of the mode of the mutation rate over time
End of explanation
"""


def mus_to_mudist(mus):
    mus_v = np.unique(mus)
    mudist = np.zeros_like(mus_v,dtype='int64')
    for i, mu in enumerate(mus_v):
        mudist[i] = np.sum(mus==mu)
    return mus_v, mudist/np.sum(mudist)

def mean_mode_mu(mus_v, mudist):
    return np.sum(mus_v*mudist)

mu_dists = OrderedDict()
for K in Ks:
    mu_dists[K]=[]
    for mu_mode in mu_modes[K]:
        mu_dists[K].append(mus_to_mudist(mu_mode[:]))

def mean_std_mu_dists(mu_dists):
    mu_dists_s = []
    for mu_dist in  mu_dists:
        mu_dists_s.append(pd.Series(mu_dist[1], mu_dist[0]))
    N = len(mu_dists_s)
    whole_index = mu_dists_s[0].index
    for i in range(1, N):
        whole_index = whole_index.union(mu_dists_s[i].index)
    for i in range(1, N):
        mu_dists_s[i]=(mu_dists_s[i]).reindex(index=whole_index, fill_value=0)
    mu_dist_total = 0
    mu_dist_total2 = 0
    for mu_dist in mu_dists_s:
        mu_dist_total = mu_dist_total + mu_dist
        mu_dist_total2 = mu_dist_total2 + mu_dist**2
    mean_mu_dist = mu_dist_total/N
    mean_squared_mu_dist = mu_dist_total2/N
    std_mu_dist = np.sqrt(mean_squared_mu_dist - mean_mu_dist**2)
    return mean_mu_dist.dropna(), std_mu_dist.dropna()/np.sqrt(N)

def bar_plot_mudist(ax, mudist_m, mudist_std):
    mus_v = mudist_m.index
    prob = mudist_m.values
    yerr = mudist_std.values
    ind = np.arange(mus_v.size)
    ax.bar(ind, prob, yerr=yerr)
    ax.set_xticks(ind);
    ax.set_xticklabels(['{:.2g}'.format(mu) for mu in mus_v], rotation=0);

"""
Explanation: The distribution of mutation rate modes as a function of population size
End of explanation
"""


for mud in mu_dists[200]:
    print(mud[0])

for mud in mu_dists[400]:
    print(mud[0])

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
bar_plot_mudist(ax, *mean_std_mu_dists([mu_dists[200][i] for i in [0,2,5,6,7,8,9,10,11,13,14,16,17]]))
ax.set_xlabel('Mode of the mutation rate, $\mu_{mode}$', fontsize=28)
ax.set_ylabel('Probability', fontsize=28)
plt.savefig('db_mu_mode_distK200.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
bar_plot_mudist(ax, *mean_std_mu_dists(mu_dists[1600]))
ax.set_xlabel('Mode of the mutation rate, $\mu_{mode}$', fontsize=28)
ax.set_ylabel('Probability', fontsize=28)
plt.savefig('db_mu_mode_distK1600.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
bar_plot_mudist(ax, *mean_std_mu_dists(mu_dists[25600]))
ax.set_xlabel('Mode of the mutation rate, $\mu_{mode}$', fontsize=28)
ax.set_ylabel('Probability', fontsize=28)
plt.savefig('db_mu_mode_distK25600.pdf')

"""
Explanation: I need to remove the runs where the distribution of mutations got screwed up because the whole population had a mutation rate of 1.
End of explanation
"""


fig = plt.figure(figsize=(18,18));
ax = fig.add_subplot(111, projection='3d');

ys, yerss = mean_std_mu_dists([mu_dists[200][i] for i in [0,2,5,6,7,8,9,10,11,13,14,16,17]])
xs = np.log2(ys.index)
ax.bar(xs,  ys, zs=10, zdir='y', alpha=.6, color=matplotlib.colors.hsv_to_rgb(np.array([0,1,1])));
colors = [matplotlib.colors.hsv_to_rgb(np.array([x,1,1])) for x in [3/32, 5/32, 8/32, 16/32, 20/32, 24/32, 28/32]]

for i, K in enumerate(Ks[1:]):
    ys, yerrs = mean_std_mu_dists(mu_dists[K])
    xs = np.log2(ys.index)
    ax.bar(xs,  ys, zs=-i*10, zdir='y', alpha=.7, color = colors[i]);
ax.set_xlabel('Mode of the mutation rate, $\mu_{mode}$ ', labelpad=40);
ax.set_ylabel('Population size, $K$', labelpad=25);
ax.set_zlabel('Probability');
ax.set_yticklabels([25600, 12800, 6400, 3200, 1600, 800, 400, 200],
                   rotation=-15, rotation_mode='anchor', ha='left', va='bottom');
ax.set_xticks(list(range(-14,-1)));
ax.set_xticklabels(['{:6.5f}'.format(i) for i in .00008*2**np.arange(14)],
                   rotation=45, rotation_mode='anchor', ha='right', va='center');
ax.plot3D(np.arange(-13,-8)+.25,10*np.arange(-6,-1),np.zeros(5), color='k',
          marker='*', markersize=20, markerfacecolor='white')
ax.plot3D(np.arange(-9,-7)+.25,10*np.arange(-2,0),np.zeros(2), color='k', marker='')
ax.plot3D(np.arange(-8,-5)+.25,10*np.arange(-1,2),np.zeros(3), color='k', marker='x', markersize=20)
plt.savefig('drift_barrier_scaling.pdf')

"""
Explanation: The mode of the mode of the mutation rate as a function of population size
To better show what's going on, I'll instead plot the empirical distribution of the dominant mutation rate for each population size. I've adjusted the axis labels manually because adjusting all the bar widths and plotting on the log scale and then attempting to rescale is hard to get looking right, manually inserting the correct labels is the easiest way I've figured out how to do things.
End of explanation
"""


fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plt.plot(f_maxes[200][0][:5*10**4],marker='')
ax.set_xlabel('generation', fontsize=28)
ax.set_ylabel('maximum fitness, $f_{max}$', fontsize=28);
plt.savefig('fitness_catastrophe.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plot_mu_trajectory(ax, mu_modes[200][0][:5*10**4])
ax.set_xlabel('generation', fontsize=28);
ax.set_ylabel('mode of the mutation rate, $\mu_{mode}$', fontsize=28);
plt.savefig('mutation_rate_during_catastrophe.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plt.plot(f_maxes[400][0][:5*10**4],marker='')
ax.set_xlabel('generation', fontsize=28)
ax.set_ylabel('maximum fitness, $f_{max}$', fontsize=28);
plt.savefig('fitness_not_catastrophe.pdf')

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111)
plot_mu_trajectory(ax, mu_modes[400][0][:5*10**4])
ax.set_xlabel('generation', fontsize=28);
ax.set_ylabel('mode of the mutation rate, $\mu_{mode}$', fontsize=28);
plt.savefig('mutation_rate_not_catastrophe.pdf')

"""
Explanation: The mutation rate landscape catastrophe
End of explanation
"""


