import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import time

eps = 0.85
L   = 10001
N   = 14
x0  = 1.0

t = np.linspace(-eps, eps, L)

def F(t, x):
    return x**2

y_tp = np.zeros((N+1, L))
y_tp[0, :] = x0

y_rt = np.zeros((N+1, L))
y_rt[0, :] = x0

i0 = L // 2
h = t[1] - t[0]

for n in range(N):
    for k in range(L):
        if k > i0:
            integrando = F(t[i0:k+1], y_tp[n, i0:k+1])
            y_tp[n+1, k] = x0 + np.trapezoid(integrando, t[i0:k+1])
        elif k < i0:
            integrando = F(t[k:i0+1], y_tp[n, k:i0+1])
            y_tp[n+1, k] = x0 - np.trapezoid(integrando, t[k:i0+1])
        else:
            y_tp[n+1, k] = x0

for n in range(N):
    for k in range(L):
        if k > i0:
            integrando = F(t[i0:k+1], y_rt[n, i0:k+1])
            y_rt[n+1, k] = x0 + np.sum(integrando[:-1])*h
        elif k < i0:
            integrando = F(t[k:i0+1], y_rt[n, k:i0+1])
            y_rt[n+1, k] = x0 - np.sum(integrando[:-1])*h
        else:
            y_rt[n+1, k] = x0

def x_exata(t):
    return 1.0 / (1.0 - t)

x_ref = x_exata(t)

print(f"{'n':<4} {'trapézio':>12} {'retângulo':>12}")
print("-" * 30)
for n in range(N+1):
    erro_tp = np.max(np.abs(y_tp[n] - x_ref))
    erro_rt = np.max(np.abs(y_rt[n] - x_ref))
    print(f"{n:<4} {erro_tp:>12.2e} {erro_rt:>12.2e}")

cores = cm.viridis(np.linspace(0, 1, N+1))
sm = cm.ScalarMappable(cmap='viridis', norm=plt.Normalize(vmin=1, vmax=N))

#  Figura 1: Trapézio 
fig1, ax1 = plt.subplots(figsize=(7, 5))
for n in range(1, N+1):
    ax1.plot(t, y_tp[n], color=cores[n])
ax1.plot(t, x_ref, linestyle='--', color='red', linewidth=2, label='solução exata')
ax1.set_title("Trapézio")
ax1.set_xlabel("$t$")
ax1.set_ylabel("$x$")
ax1.set_xlim(-eps, eps)
ax1.set_ylim(0, 5)
ax1.legend()
ax1.grid(True)
plt.colorbar(sm, ax=ax1, label='iterado $n$')
plt.tight_layout()

plt.show()

#  Figura 2: Retângulo 
fig2, ax2 = plt.subplots(figsize=(7, 5))
for n in range(1, N+1):
    ax2.plot(t, y_rt[n], color=cores[n])
ax2.plot(t, x_ref, linestyle='--', color='red', linewidth=2, label='solução exata')
ax2.set_title("Retângulo")
ax2.set_xlabel("$t$")
ax2.set_ylabel("$x$")
ax2.set_xlim(-eps, eps)
ax2.set_ylim(0, 5)
ax2.legend()
ax2.grid(True)
sm2 = cm.ScalarMappable(cmap='viridis', norm=plt.Normalize(vmin=1, vmax=N))
plt.colorbar(sm2, ax=ax2, label='iterado $n$')
plt.tight_layout()

plt.show()

#  Figura 3: Erro pontual 
fig3, ax3 = plt.subplots(figsize=(7, 5))
ax3.plot(t, np.abs(y_tp[N] - x_ref), label="erro trapézio")
ax3.plot(t, np.abs(y_rt[N] - x_ref), label="erro retângulo")
ax3.set_yscale('log')
ax3.set_xlabel("$t$")
ax3.set_ylabel("erro pontual")
ax3.legend()
ax3.grid(True)
plt.tight_layout()

plt.show()