import matplotlib.pyplot as plt
import numpy as np
from io import StringIO
from math import lgamma

# Left: variance of the mean of n estimators (variance 1, pairwise correlation rho),
# formula rho + (1 - rho)/n against a simulation.
# Right: probability that the majority of n voters is right, each voter right with p = 0.6,
# when each voter copies one shared opinion with probability c (correlation c² between voters).
rng = np.random.default_rng(0)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9.5, 3.6))

colors = {0.0: '#3fb950', 0.1: '#58a6ff', 0.3: '#f0883e', 0.6: '#ff7b72'}
n_line = np.unique(np.logspace(0, 3, 200).astype(int))
n_dots = [1, 3, 10, 30, 100, 300, 1000]
for rho, col in colors.items():
    ax1.plot(n_line, rho + (1 - rho) / n_line, color=col, lw=1.8, label=f'ρ = {rho:g}')
    sim = []
    for n in n_dots:
        est = np.sqrt(rho) * rng.normal(size=(4000, 1)) + np.sqrt(1 - rho) * rng.normal(size=(4000, n))
        sim.append(est.mean(axis=1).var())
    ax1.plot(n_dots, sim, 'o', color=col, ms=4)
    if rho > 0:
        ax1.axhline(rho, color=col, ls=':', lw=1)
        ax1.text(1000, rho - 0.015, f'floor ρσ² = {rho:g}', color=col, fontsize=7.5, ha='right', va='top')
ax1.set_xscale('log')
ax1.set_xlabel('number of estimators averaged, n')
ax1.set_ylabel('variance of the average (σ² = 1)')
ax1.set_ylim(0, 1.05)
ax1.legend(fontsize=7.5, loc='upper right', title='correlation', title_fontsize=7.5, bbox_to_anchor=(1, 1.0))
ax1.set_title('Averaging: lines = ρσ² + (1 − ρ)σ²/n, dots = simulation', fontsize=9, loc='left')


def pmf(n, p):
    """Binomial probabilities P(K = 0..n) for K ~ Bin(n, p), computed in logs."""
    k = np.arange(n + 1)
    logc = np.array([lgamma(n + 1) - lgamma(i + 1) - lgamma(n - i + 1) for i in k])
    with np.errstate(divide='ignore'):
        return np.exp(logc + k * np.log(p) + (n - k) * np.log1p(-p)) if 0 < p < 1 else (k == (n if p else 0)).astype(float)


def at_least(m, n, p):
    """P(Bin(n, p) >= m)."""
    return pmf(n, p)[max(m, 0):].sum() if m <= n else 0.0


def majority(n, p, c):
    need = n // 2 + 1
    wk = pmf(n, c)                              # P(k voters copy the shared opinion)
    right_if_shared_right = np.array([at_least(need - k, n - k, p) for k in range(n + 1)])
    right_if_shared_wrong = np.array([at_least(need, n - k, p) for k in range(n + 1)])
    return p * (wk * right_if_shared_right).sum() + (1 - p) * (wk * right_if_shared_wrong).sum()


n_odd = np.unique(np.r_[1, 3, 5, 7, 9, (np.logspace(1, 3, 60).astype(int) // 2) * 2 + 1])
for c, col in [(0.0, '#3fb950'), (0.1, '#58a6ff'), (0.3, '#f0883e')]:
    acc = [majority(n, 0.6, c) for n in n_odd]
    ax2.plot(n_odd, acc, color=col, lw=1.8, label=f'copy prob. c = {c:g}  (ρ = {c * c:.2g})')
    a101 = majority(101, 0.6, c)
    ax2.plot(101, a101, 'o', color=col, ms=4)
    ax2.annotate(f'{a101:.2f}', (101, a101), xytext=(4, -9 if c < 0.2 else 4), textcoords='offset points',
                 color=col, fontsize=7.5)
ax2.axhline(0.6, color='gray', ls=':', lw=1)
ax2.text(1.1, 0.585, 'one voter: 0.6', color='gray', fontsize=7.5, va='top')
ax2.set_xscale('log')
ax2.set_xlim(0.9, 1200)
ax2.set_ylim(0.5, 1.02)
ax2.set_xlabel('number of voters, n (odd)')
ax2.set_ylabel('P(majority is right)')
ax2.legend(fontsize=7.5, loc='upper left')
ax2.set_title('Majority vote, each voter right 60% of the time', fontsize=9, loc='left')
for ax in (ax1, ax2):
    ax.tick_params(labelsize=8)
    ax.set_xticks([1, 10, 100, 1000], ['1', '10', '100', '1000'])

plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()