# Archivo complementario del capítulo 23.
# Reproduce las seis figuras del capítulo y ejecuta las comprobaciones.
# Los tres primeros bloques son idénticos a los del texto.
import matplotlib
matplotlib.use("Agg")
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

plt.rcParams.update({"font.size": 9, "axes.spines.top": False,
                     "axes.spines.right": False})
coma = FuncFormatter(lambda v, _: f"{v:g}".replace(".", ","))


def simular(alpha, beta, inicial=(100.0, 0.0), pasos=30):
    a, b = float(alpha), float(beta)
    if not (np.isfinite(a) and np.isfinite(b)):
        raise ValueError("Las fracciones deben ser finitas.")
    if not (0.0 <= a <= 1.0 and 0.0 <= b <= 1.0):
        raise ValueError("Las fracciones deben estar entre 0 y 1.")
    if (isinstance(pasos, (bool, np.bool_))
            or not isinstance(pasos, (int, np.integer))
            or pasos < 0):
        raise ValueError("pasos debe ser un entero no negativo.")
    z0 = np.array(inicial, dtype=float, copy=True)
    if z0.shape != (2,) or not np.all(np.isfinite(z0)):
        raise ValueError("Se requieren dos cantidades finitas.")
    if np.any(z0 < 0):
        raise ValueError("Las cantidades deben ser no negativas.")
    matriz = np.array([[1.0-a, b], [a, 1.0-b]])
    historia = np.empty((pasos + 1, 2), dtype=float)
    historia[0] = z0
    with np.errstate(over="raise", invalid="raise"):
        for n in range(pasos):
            historia[n+1] = matriz @ historia[n]
    return historia


base = simular(0.20, 0.10)
lenta = simular(0.10, 0.05)
rondas = np.arange(base.shape[0])
x_equilibrio = 100.0 / 3.0
salida = Path("figuras")
salida.mkdir(exist_ok=True)

copias = {"base": base.copy(), "lenta": lenta.copy()}

fig, ax = plt.subplots(figsize=(6.4, 3.0))
ax.plot(rondas, base[:, 0], marker="o", markersize=3, color="black",
        label="A (α = 0,20; β = 0,10)")
ax.plot(rondas, lenta[:, 0], marker="s", markersize=3, color="0.45",
        linestyle="--", label="B (α = 0,10; β = 0,05)")
ax.axhline(x_equilibrio, linestyle=":", color="0.4",
           label="Equilibrio")
ax.set(xlabel="Ronda n", ylabel="Cantidad x (Q)",
       xlim=(0, 30), ylim=(0, 105))
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=8, frameon=False)
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap23_fig1.{ext}", dpi=300,
                bbox_inches="tight")
plt.close(fig)

s = np.linspace(0.0, 1.5, 151)
x_5 = np.empty(s.size)
x_30 = np.empty(s.size)
for j, valor in enumerate(s):
    registro = simular(2.0 * valor / 3.0, valor / 3.0)
    x_5[j] = registro[5, 0]
    x_30[j] = registro[30, 0]
referencia_30 = 100.0 / 3.0 + (200.0 / 3.0) * (1.0-s)**30
if not np.allclose(x_30, referencia_30,
                   rtol=1e-12, atol=1e-10):
    raise RuntimeError("Discrepancia en el barrido paramétrico.")

# Figura 23.2. Ampliación de las rondas 15 a 30 del caso A.
fig, ax = plt.subplots(figsize=(6.4, 3.0))
ax.plot(rondas[15:], base[15:, 0], marker="o", markersize=3, color="black",
        label="A (α = 0,20; β = 0,10)")
ax.axhline(x_equilibrio, linestyle=":", color="0.4", label="Equilibrio")
ax.set(xlabel="Ronda n", ylabel="Cantidad x (Q)", xlim=(15, 30),
       ylim=(33.32, 33.68))
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=8, frameon=False)
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap23_fig2.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Figura 23.3. Desviación normalizada en escala logarítmica.
eta_a = np.abs(base[:, 0] - x_equilibrio) / 100.0
eta_b = np.abs(lenta[:, 0] - x_equilibrio) / 100.0
fig, ax = plt.subplots(figsize=(6.4, 3.0))
ax.semilogy(rondas, eta_a, marker="o", markersize=3, color="black",
            nonpositive="mask", label="A")
ax.semilogy(rondas, eta_b, marker="s", markersize=3, color="0.45",
            linestyle="--", nonpositive="mask", label="B")
ax.axhline(0.01, linestyle=":", color="0.4", label="0,01 (desviación de 1 Q)")
for n0, eta, dx in ((12, eta_a, -1.2), (26, eta_b, 0.6)):
    ax.plot([n0], [eta[n0]], marker="o", markersize=8, mfc="none",
            color="black", linestyle="None")
    ax.annotate(f"ronda {n0}", xy=(n0, eta[n0]), xytext=(n0 + dx, eta[n0] / 5),
                fontsize=8, ha="center")
ax.set(xlabel="Ronda n", ylabel=r"$|x - x_{\mathrm{eq}}|\,/\,100$", xlim=(0, 30))
ax.legend(fontsize=8, frameon=False)
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap23_fig3.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Figura 23.4. Plano de estados, rondas 0 a 10.
fig, ax = plt.subplots(figsize=(4.6, 4.4))
recta = np.array([0.0, 100.0])
ax.plot(recta, 100.0 - recta, color="0.75", linewidth=0.8, label="x + y = 100")
ax.plot(base[:11, 0], base[:11, 1], linestyle="None", marker="o",
        markersize=4.5, mfc="none", color="black", label="A")
ax.plot(lenta[:11, 0], lenta[:11, 1], linestyle="None", marker="s",
        markersize=3.5, color="0.45", label="B")
ax.plot([x_equilibrio], [100.0 - x_equilibrio], linestyle="None", marker="*",
        markersize=10, color="black", label="Equilibrio")
ax.annotate("A, ronda 5", xy=(base[5, 0], base[5, 1]), xytext=(12, 42),
            fontsize=8, arrowprops=dict(arrowstyle="-", color="0.5"))
ax.annotate("B, ronda 5", xy=(lenta[5, 0], lenta[5, 1]), xytext=(72, 58),
            fontsize=8, arrowprops=dict(arrowstyle="-", color="0.5"))
ax.set(xlabel="x (Q)", ylabel="y (Q)", xlim=(0, 105), ylim=(0, 105))
ax.set_aspect("equal", adjustable="box")
ax.xaxis.set_major_formatter(coma)
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=7, frameon=False, loc="upper right")
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap23_fig4.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Figura 23.5. Alternancia y submuestreo en rondas pares.
alternante = simular(0.8, 0.6, pasos=12)
r = np.arange(alternante.shape[0])
fig, ax = plt.subplots(figsize=(6.4, 3.0))
ax.plot(r, alternante[:, 0], marker="o", markersize=3, color="black",
        label="Todas las rondas")
ax.plot(r[::2], alternante[::2, 0], linestyle="None", marker="s",
        markersize=7, mfc="none", color="0.35", label="Rondas pares")
ax.axhline(300.0 / 7.0, linestyle=":", color="0.4",
           label="Equilibrio (300/7)")
ax.set(xlabel="Ronda n", ylabel="Cantidad x (Q)", xlim=(0, 12), ylim=(0, 105))
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=8, frameon=False)
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap23_fig5.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Figura 23.6. Estados después de 5 y de 30 rondas según el parámetro s.
fig, ax = plt.subplots(figsize=(6.4, 3.0))
ax.plot(s, x_5, marker="o", markersize=2.5, markevery=5, color="black",
        label="Después de 5 rondas")
ax.plot(s, x_30, marker="s", markersize=2.5, markevery=5, color="0.55",
        label="Después de 30 rondas")
positivos = s > 0
ax.plot(s[positivos], np.full(positivos.sum(), x_equilibrio), linestyle="--",
        color="0.3", label="Equilibrio (s > 0)")
ax.set(xlabel="Parámetro s", ylabel="Cantidad x (Q)", xlim=(0, 1.5),
       ylim=(0, 105))
ax.xaxis.set_major_formatter(coma)
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=8, frameon=False)
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap23_fig6.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Comprobaciones.
for datos, rho in ((base, 0.7), (lenta, 0.85)):
    if not np.allclose(datos.sum(axis=1), 100.0, rtol=1e-12, atol=1e-10):
        raise RuntimeError("El total no se conserva.")
    referencia = x_equilibrio + (100.0 - x_equilibrio) * rho**rondas
    if not np.allclose(datos[:, 0], referencia, rtol=1e-12, atol=1e-10):
        raise RuntimeError("Discrepancia con la solución analítica.")
if not (np.array_equal(base, copias["base"])
        and np.array_equal(lenta, copias["lenta"])):
    raise RuntimeError("Dibujar modificó los arreglos recibidos.")
if not np.allclose(alternante[:7, 0],
                   [100, 20, 52, 39.2, 44.32, 42.272, 43.0912]):
    raise RuntimeError("Discrepancia en la trayectoria alternante.")
if (int(np.argmax(eta_a <= 0.01)), int(np.argmax(eta_b <= 0.01))) != (12, 26):
    raise RuntimeError("Las rondas de la tolerancia no coinciden con el texto.")
print("Seis figuras generadas y comprobaciones satisfactorias.")