# %%
import numpy as np
from matplotlib import pyplot as plt


def p_win(p_ov):
    return np.exp(1 / (1 - p_ov) * np.log(p_ov)) + 1 - np.exp(p_ov / (1 - p_ov) * np.log(p_ov))


p_ovs = np.logspace(-5, 0, 1000)
plt.figure(figsize=(4, 3), dpi=300)
plt.plot(p_ovs, p_win(p_ovs), label="indistinguished")
plt.plot(p_ovs, p_ovs, label="distinguished", c="k")
plt.xscale("log")
plt.xlabel("p_ov")
plt.ylabel("P(humanity wins)")
plt.title("P(humanity wins) if e_ov = 1")
plt.legend()
plt.show()
# %%
plt.figure(figsize=(4, 3), dpi=300)
p_ovs = np.logspace(-5, 0, 1000)
colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
for s, c in zip([1, 0.99, 0.9, 0.5, 0.1, 0.01], colors):
    p_ovs_p = p_ovs / (p_ovs + s * (1 - p_ovs))
    plt.plot(p_ovs, p_win(p_ovs_p), label=f"s={s}", color=c)
    if s != 1:
        plt.plot(p_ovs, 1 - (1 - p_win(p_ovs)) * s, label=f"s_loud={s}", color=c, linestyle="--")
plt.xlabel("p_ov")
plt.ylabel("P(humanity wins)")
plt.xscale("log")
plt.title(f"P(humanity wins) if e_ov = 1\nGiven silent or loud failures")
plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
# %%
plt.figure(figsize=(3, 3), dpi=300)
p_ovs = np.logspace(-5, 0, 1000)
colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
for s, c in zip([1, 0.99, 0.9, 0.5, 0.1, 0.01], colors):
    p_ovs_p = p_ovs / (p_ovs + s * (1 - p_ovs))
    lamb = np.log(1/p_ovs_p)/(1 - p_ovs_p) / (p_ovs + s * (1 - p_ovs))
    plt.plot(p_ovs, lamb, label=f"s={s}", color=c)
plt.axhline(1, color="k", linestyle="--", label="1 attempt\non average")
plt.xlabel("p_ov")
plt.ylabel("λ T")
plt.xscale("log")
plt.yscale("log")
plt.title(f"Average number of take-over attempts\n(λ T)")
plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
# %%
# combination of the two plots above
fig, axs = plt.subplots(1, 2, figsize=(7, 3), dpi=300)
p_ovs = np.logspace(-5, 0, 1000)
colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
ax = axs[0]
for s, c in zip([1, 0.99, 0.9, 0.5, 0.1, 0.01], colors):
    p_ovs_p = p_ovs / (p_ovs + s * (1 - p_ovs))
    ax.plot(p_ovs, p_win(p_ovs_p), color=c)
ax.set_xlabel("p_ov")
ax.set_ylabel("P(humanity wins)")
ax.set_xscale("log")
ax.set_title(f"P(humanity wins)")
ax = axs[1]
for s, c in zip([1, 0.99, 0.9, 0.5, 0.1, 0.01], colors):
    p_ovs_p = p_ovs / (p_ovs + s * (1 - p_ovs))
    lamb = np.log(1/p_ovs_p)/(1 - p_ovs_p) / (p_ovs + s * (1 - p_ovs))
    ax.plot(p_ovs, lamb, label=f"s={s}", color=c)
ax.axhline(1, color="k", linestyle="--")
ax.set_xlabel("p_ov")
ax.set_ylabel("λ T")
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_title(f"Avg nb of take-over attempts")
ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
plt.tight_layout()

# %%
from scipy.optimize import bisect


def p_win_mult(Tis, sis, cis, B):
    n_tasks = len(Tis)
    assert n_tasks == len(sis) == len(cis) == len(Tis)

    def over_budget(P):
        return B - sum([Ti * ci * P / (P + (1 - P) / si) for Ti, si, ci in zip(Tis, sis, cis)])

    P = bisect(over_budget, 0, 1, maxiter=1000)
    Lambda = 1 / n_tasks * np.log(1 / P) / (1 - P)
    return np.exp(-Lambda * n_tasks) + 1 - np.exp(-Lambda * n_tasks * P)


def close_to(a, b):
    return np.abs(a - b) < 1e-5


# assert close_to(p_win_mult([2], [0.5], [4], 3), 0.5)
Tis = [1, 1]
sis = [0.1, 0.1]
cis = [4, 4]
B = 3
p_ov = B / (cis[0] * Tis[0]) / 2
p_ov_ = p_ov / (p_ov + sis[0] * (1 - p_ov))
print(p_ov, p_ov_)
old = p_win(p_ov_)
new = p_win_mult(Tis, sis, cis, B)
assert close_to(old, new), f"{old} != {new}"
# %%
plt.figure(figsize=(5, 3), dpi=300)
Ttot = 1
s0 = 1
n_tasks = 2
cis = [1, 1]
kept_colors = [colors[0]] + colors[3:]
for s1, c in zip([1, 0.5, 0.1, 0.01], kept_colors):
    for T1, ls in zip([0.9, 0.99, 1], ["dotted", "--", "-"]):
        if s1 == 1 and T1 != 1:
            continue

        p_ovs = np.logspace(-5, 0, 1000)
        Bs = Ttot * cis[0] * p_ovs
        p_wins = [p_win_mult([Ttot - T1, T1], [s0, s1], cis, B) for B in Bs]
        plt.plot(p_ovs, p_wins, label=f"s={s1} r={T1}", color=c, linestyle=ls)
plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
plt.xlabel("p_ov")
plt.ylabel("P(humanity wins)")
plt.xscale("log")
plt.title(
    f"P(humanity wins)\nif a fraction r of the episodes are 'low danger'\nand have a takeover success rate s.\nThe other episodes have a success rate of 1"
)
# %%
plt.figure(figsize=(6, 3), dpi=300)
p_ovs = np.logspace(-5, 0, 1000)
colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
for s, c in zip([1, 0.99, 0.9, 0.5, 0.1, 0.01], colors):
    p_wins = 1 - 1 / (1 + p_ovs / (s * (1 - p_ovs)))
    plt.plot(p_ovs, p_wins, label=f"s={s}", color=c)
plt.xlabel("p_ov")
plt.ylabel("P(humanity wins)")
plt.xscale("log")
plt.title(f"P(humanity wins) if e_ov = 1\nIn an infinite game")
plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
# %%
from scipy.optimize import fsolve
import numpy as np
import sympy as sp


def p_win_given_lambda(lamb, k, p, T):
    return 1 - (1 - (1 - np.power(lamb, k)) ** (T * (1 - p))) * np.power(1 - lamb, T * p)


# lambda_val_, k_, T_, p_ = sp.symbols("lambda k T p")
# der_p_win_given_lambda = sp.diff(p_win_given_lambda(lambda_val_, k_, p_, T_), lambda_val_)


def der_p_win_given_lambda(lamb, k, p, T):
    return -T * k * lamb**k * (1 - lamb) ** (T * p) * (1 - lamb**k) ** (T * (1 - p)) * (1 - p) / (
        lamb * (1 - lamb**k)
    ) + T * p * (1 - lamb) ** (T * p) * (1 - (1 - lamb**k) ** (T * (1 - p))) / (1 - lamb)


def other_zero_expression(lamb, k, p, T):
    # return -T * k * lamb**k * (1 - lamb) ** (T * p) * (1 - lamb**k) ** (T * (1 - p)) * (1 - p) * (
    #     1 - lamb
    # ) + T * p * (1 - lamb) ** (T * p) * (1 - (1 - lamb**k) ** (T * (1 - p))) * (lamb * (1 - lamb**k))
    t1 = np.log(k) + k*np.log(lamb) + T*p*np.log(1 - lamb) + T*(1 - p)*np.log(1 - lamb**k) + np.log(1 - p) + np.log(1 - lamb)
    t2 = np.log(p) + T*p*np.log(1 - lamb) + np.log(1 - (1 - lamb**k) ** (T * (1 - p))) + np.log(lamb) + np.log(1 - lamb**k)
    return t2 - t1

# lambdas = np.linspace(0, 1, 1000)
# T = 100
# p = 0.001
# k = 2
# plt.plot(lambdas, p_win_given_lambda(lambdas, k, p, T))
# plt.plot(lambdas, der_p_win_given_lambda(lambdas, k, p, T) > 0)
# plt.plot(lambdas, other_zero_expression(lambdas, k, p, T) > 0)
# %%

def coordination_p_win(k, p, T):
    def func(lamb):
        return other_zero_expression(lamb, k, p, T)
    lambd = bisect(func, 1e-6, 1-1e-6, maxiter=1000)
    return p_win_given_lambda(lambd, k, p, T)


coordination_p_win(3, 0.9, 1000)

plt.figure(figsize=(6, 4), dpi=300)
for T, c in zip([10, 100, 1000, 10000], colors):
    for k, ls in zip([1, 2, 3, 6], ["-", "--", "-.", "dotted"]):
        if k == 1 and T != 10000:
            continue
        N = 1
        p_ovs = np.logspace(-5, 0, 100)
        p_wins = [coordination_p_win(k, p, T) for p in p_ovs]
        plt.plot(p_ovs, p_wins, label=f"k={k} T/N={T}", color=c, linestyle=ls)
plt.xlabel("p_ov")
plt.ylabel("P(humanity wins)")
plt.xscale("log")
plt.title(f"P(humanity wins)\nin a cooperative take-over scenario\nwith k AIs")
plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))

# %%