from functools import cache
import numpy as np
import matplotlib.pyplot as plt
import itertools
import math
from decimal import Decimal, getcontext

def make_pq(p1, p2, p3):
    @cache
    def p(i, k):
        if k < 0:
            raise ValueError("p called with k < 0")
        if k == 0:
            return [p1, p2, p3][i - 1]
        return p(i, k - 1) * (p(i, k - 1) + 2 * p(i % 3 + 1, k - 1))
    
    @cache
    def q(i, k):
        if k < 0:
            raise ValueError("q called with k < 0")
        if k == 0:
            return p(i, 0) * (p(i, 0) + 2 * p((i + 1) % 3 + 1, 0))
        return (
            q(i, k - 1) * p(i%3 + 1, k - 1) * (p(i%3 + 1, k - 1) + 2 * p(i, k - 1)) +
            q(i, k - 1) * p(i, k - 1) * (p(i, k - 1) + 2 * p((i + 1)%3 + 1, k - 1)) +
            q(i%3 + 1, k - 1) * p(i, k - 1) * (p(i, k - 1) + 2 * p((i + 1)%3 + 1, k - 1))
        )
    
    return (p, q)

p, q = make_pq(.1, .8, .1)

fig, ax = plt.subplots()
for i in [1, 2, 3]:
    data = []
    for k in range(50):
        data.append(p(i, k))
#         if data[-1] > 1:
#             print(i, k)
#             error
#     print(data[30:40])
    ax.plot(data, label=f"$i={i}$")
ax.legend()
plt.show()

fig, ax = plt.subplots()
for i in [1, 2, 3]:
    data = []
    for k in range(50):
        data.append(q(i, k))
    ax.plot(data, label=f"$i={i}$")
ax.legend()
plt.show()

def plot(p1, p2, p3, k, label=None):
    fig, ax = plt.subplots()
    p, q = make_pq(p1, p2, p3)
    colors = ["red", "blue", "green", "pink", "aqua", "lime"]
    color_index = 0
    for fun in [p, q]:
        for i in [1, 2, 3]:
            data = []
            for k_val in range(k):
                data.append(fun(i, k_val))
            ax.plot(data, label=f"${fun.__name__}_{i}$", color=colors[color_index])
            color_index += 1
    ax.legend()
    ax.set_title(f"$p_1^{(0)}={p1}$, $p_2^{(0)}={p2}$, $p_3^{(0)}={p3}$, $k \\in [0, {k}]$" if label is None else label)
    plt.savefig(f"./de/p1_{p1}-p2_{p2}-p3_{p3}-k_{k}.png")
    plt.show()

plot(.4, .2, .4, 50)
plot(.3, .4, .3, 50)
plot(.2, .6, .2, 50)
plot(.1, .8, .1, 50)
plot(.05, .9, .05, 50)
plot(.025, .95, .025, 50)

getcontext().prec = 1500
delta = 10
plot(Decimal(delta) / Decimal(100), Decimal(100 - 2 * delta) / Decimal(100), Decimal(delta) / Decimal(100), 4096)

getcontext().prec = 3000
delta = 10
plot(Decimal(delta) / Decimal(100), Decimal(100 - 2 * delta) / Decimal(100), Decimal(delta) / Decimal(100), 8192)

getcontext().prec = 1500
epsilon = Decimal(1) / Decimal(2 ** 75)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 4096, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-75}$, $k \in [0, 4096]$")

getcontext().prec = 3000
epsilon = Decimal(1) / Decimal(2 ** 80)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 8192, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-80}$, $k \in [0, 8192]$")

getcontext().prec = 6000
epsilon = Decimal(1) / Decimal(2 ** 80)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 8192, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-80}$, $k \in [0, 8192]$")

getcontext().prec = 6000
epsilon = Decimal(1) / Decimal(2 ** 75)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 8192, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-75}$, $k \in [0, 8192]$")

getcontext().prec = 12000
epsilon = Decimal(1) / Decimal(2 ** 75)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 8192, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-75}$, $k \in [0, 8192]$")

getcontext().prec = 12000
epsilon = Decimal(1) / Decimal(2 ** 15)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 8192, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-15}$, $k \in [0, 8192]$")

getcontext().prec = 12000
epsilon = Decimal(1) / Decimal(2 ** 18)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 8192, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-18}$, $k \in [0, 8192]$")

getcontext().prec = 12000
epsilon = Decimal(1) / Decimal(2 ** 20)
plot(epsilon, Decimal(1) - 2*epsilon, epsilon, 8192, label="$p_1^{(0)} = p_3^{(0)} = \\epsilon = 2^{-20}$, $k \in [0, 8192]$")