# Archivo complementario del capítulo 30.
# Reproduce las cuatro figuras, los cuadros de escenarios y de convergencia
# y el resumen de verificación del capítulo.
import math
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
import scipy
from scipy.integrate import solve_ivp

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

# Especificación del modelo base (idéntica al bloque del capítulo).
lam = 8.0
k = 0.10
W0 = 120.0
H = 90.0


def rhs(t, y):
    return [lam - k*y[0]]


def evento(t, y):
    return y[0] - H


evento.terminal = True
evento.direction = -1

sol = solve_ivp(rhs, (0.0, 50.0), [W0], rtol=1e-10, atol=1e-12,
                events=evento, dense_output=True)

# Solución exacta y tiempo del evento.
def W_exacta(t, kk=k, lam=lam, W0=W0):
    return lam/kk + (W0 - lam/kk)*np.exp(-kk*np.asarray(t))


def t_cruce(kk, lam=lam, W0=W0, H=H):
    weq = lam/kk
    if not weq < H < W0:
        return math.nan
    return -math.log((H - weq)/(W0 - weq))/kk


t_num = sol.t_events[0][0]
t_ex = math.log(4.0)/0.1

# Verificaciones.
assert sol.y[0, 0] == W0
assert abs(rhs(0.0, [lam/k])[0]) < 1e-12
assert np.all(np.diff(sol.y[0]) < 0)
assert np.all(sol.y[0] > 0)
assert np.allclose(sol.y[0], W_exacta(sol.t), rtol=1e-9, atol=1e-9)
assert abs(t_num - t_ex) < 1e-8

# Cuadro de escenarios.
escenarios = [(kk, lam/kk, t_cruce(kk))
              for kk in (0.06, 0.08, 0.09, 0.095, 0.10, 0.12, 0.15)]

# Convergencia de Euler en t = 10.
W10 = float(W_exacta(10.0))
convergencia = []
for h in (2.0, 1.0, 0.5, 0.25, 0.125):
    W = W0
    for _ in range(int(round(10.0/h))):
        W = W + h*(lam - k*W)
    convergencia.append((h, W, abs(W - W10)))
errores = [c[2] for c in convergencia]
assert all(e1 > e2 for e1, e2 in zip(errores, errores[1:]))

# Figura 30.1.
t = np.linspace(0.0, 50.0, 500)
fig, ax = plt.subplots(figsize=(6.4, 3.1))
ax.plot(t, W_exacta(t), color="0.7", lw=2.6, label="Solución analítica")
ax.plot(sol.t, sol.y[0], ls="None", marker="o", ms=3, color="black",
        label="Pasos del integrador")
ax.axhline(lam/k, color="0.4", ls=":", lw=1.0, label="Equilibrio, W = 80")
ax.axhline(H, color="0.4", ls="--", lw=1.0, label="Umbral, H = 90")
ax.plot([t_num], [H], marker="D", ms=6, color="black", ls="None",
        label="Evento, t = 13,86")
ax.set(xlabel="Tiempo (días)", ylabel="Carga W", xlim=(0, 50), ylim=(70, 125))
ax.xaxis.set_major_formatter(coma)
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=7.5, frameon=False, loc="upper right")
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap30_fig1.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Figura 30.2.
t = np.linspace(0.0, 60.0, 600)
fig, ax = plt.subplots(figsize=(6.4, 3.1))
estilos = [("-", "0.6"), ("--", "0.6"), (":", "black"), ("-", "black"),
           ("--", "black"), ("-.", "black")]
for kk, (ls, col) in zip((0.06, 0.08, 0.09, 0.10, 0.12, 0.15), estilos):
    ax.plot(t, W_exacta(t, kk), ls=ls, color=col, lw=1.3,
            label=f"k = {kk:.2f}".replace(".", ","))
ax.axhline(H, color="0.5", ls=":", lw=1.0)
ax.text(60.5, H, "H", va="center", fontsize=8)
ax.set(xlabel="Tiempo (días)", ylabel="Carga W", xlim=(0, 60), ylim=(40, 140))
ax.xaxis.set_major_formatter(coma)
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=7.5, frameon=False, ncol=3, loc="upper right")
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap30_fig2.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Figura 30.3.
kc = lam/H
kk = np.linspace(kc + 1e-5, 0.20, 600)
fig, ax = plt.subplots(figsize=(6.4, 3.1))
ax.plot(kk, [t_cruce(v) for v in kk], color="black", lw=1.4,
        label="Tiempo exacto hasta H")
kp = [e[0] for e in escenarios if not math.isnan(e[2])]
tp = [e[2] for e in escenarios if not math.isnan(e[2])]
ax.plot(kp, tp, ls="None", marker="o", ms=4.5, mfc="white", color="black",
        label="Valores del cuadro 30.2")
ax.axvline(kc, color="0.5", ls="--", lw=1.0, label="k crítico = 8/90")
ax.set(xlabel="Tasa de salida k (día⁻¹)", ylabel="Tiempo hasta H (días)",
       xlim=(0.08, 0.20), ylim=(0, 80))
ax.xaxis.set_major_formatter(coma)
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=7.5, frameon=False, loc="upper right")
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap30_fig3.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Figura 30.4.
t = np.linspace(0.0, 35.0, 500)
fig, ax = plt.subplots(figsize=(6.4, 3.1))
ax.plot(t, W_exacta(t), color="black", lw=1.4, label="Salida proporcional, k = 0,10")
ax.plot(t, np.maximum(W0 - 4.0*t, 0.0), color="0.5", lw=1.4, ls="--",
        label="Capacidad fija de 12 por día")
ax.axhline(H, color="0.6", ls=":", lw=1.0)
ax.plot([t_ex, 7.5], [H, H], ls="None", marker="D", ms=5, color="black")
ax.annotate("13,86 días", xy=(t_ex, H), xytext=(t_ex + 2, H + 12), fontsize=8,
            arrowprops=dict(arrowstyle="-", color="0.5"))
ax.annotate("7,5 días", xy=(7.5, H), xytext=(1.0, 60), fontsize=8,
            arrowprops=dict(arrowstyle="-", color="0.5"))
ax.set(xlabel="Tiempo (días)", ylabel="Carga W", xlim=(0, 35), ylim=(0, 125))
ax.xaxis.set_major_formatter(coma)
ax.yaxis.set_major_formatter(coma)
ax.legend(fontsize=7.5, frameon=False, loc="upper right")
fig.tight_layout()
for ext in ("png", "svg"):
    fig.savefig(salida / f"cap30_fig4.{ext}", dpi=300, bbox_inches="tight")
plt.close(fig)

# Resumen.
print(f"SciPy {scipy.__version__}")
print(f"Tiempo numérico {t_num:.10f}, exacto {t_ex:.10f}, "
      f"diferencia {t_num - t_ex:.1e}, evaluaciones {sol.nfev}")
for kk, weq, tc in escenarios:
    print(f"k = {kk:.3f}  equilibrio {weq:.2f}  tiempo "
          + ("sin cruce" if math.isnan(tc) else f"{tc:.2f}"))
for h, W, e in convergencia:
    print(f"Euler h = {h:<6} W(10) = {W:.6f}  error {e:.6f}")
print("Verificaciones satisfactorias.")