import torch
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import gamma  # complex gamma
import mpmath  # for ground truth comparison

# ------------------------------------------------------------
# 1. Build the frozen wave shaper (random complex exponentials)
# ------------------------------------------------------------
K = 400                         # number of basis functions
omega = torch.randn(K) * 8.0    # random frequencies (fixed)
omega.requires_grad_(False)

# Learnable coefficients (complex)
c = torch.nn.Parameter(torch.randn(K, dtype=torch.complex64) * 0.1)

def zeta_approx(t):
    """
    t: real tensor (imaginary part of s = 0.5 + i*t)
    returns ζ(0.5 + i*t) as complex tensor
    """
    # phi(t) = exp(i * omega_k * t)   shape (len(t), K)
    phase = torch.outer(t, omega)           # (len(t), K)
    phi = torch.exp(1j * phase)
    return phi @ c                           # linear superposition

# ------------------------------------------------------------
# 2. Functional equation residual (using SciPy gamma)
# ------------------------------------------------------------
def functional_residual(t_points):
    """
    t_points: 1D tensor of real collocation points
    """
    t = t_points.numpy()
    s = 0.5 + 1j * t
    # Compute ζ(0.5+it) and ζ(0.5-it) from our model
    zeta_s = zeta_approx(t_points).detach().numpy()
    zeta_1_minus_s = zeta_approx(-t_points).detach().numpy()

    # Gamma factor: 2^s * π^{s-1} * sin(π s/2) * Γ(1-s)
    pow2 = 2.0 ** s
    pow_pi = np.pi ** (s - 1.0)
    sin_term = np.sin(np.pi * s / 2.0)
    gamma_term = gamma(1.0 - s)   # complex gamma
    rhs = pow2 * pow_pi * sin_term * gamma_term * zeta_1_minus_s

    # Convert back to torch
    rhs_torch = torch.tensor(rhs, dtype=torch.complex64)
    zeta_s_torch = zeta_approx(t_points)
    return torch.mean(torch.abs(zeta_s_torch - rhs_torch) ** 2)

# ------------------------------------------------------------
# 3. Known value at t=0
# ------------------------------------------------------------
t0 = torch.tensor([0.0])
zeta_half_true = -1.4603545088095868  # ζ(0.5)
loss_known = torch.abs(zeta_approx(t0) - zeta_half_true) ** 2

# ------------------------------------------------------------
# 4. Optimisation
# ------------------------------------------------------------
t_colloc = torch.linspace(1.0, 20.0, 150)   # collocation points (avoid t=0 singularity? t=0 is fine)
optimizer = optim.Adam([c], lr=0.03)

for step in range(3000):
    optimizer.zero_grad()
    loss_res = functional_residual(t_colloc)
    loss = loss_res + 20.0 * loss_known
    loss.backward()
    optimizer.step()
    if step % 500 == 0:
        print(f"Step {step}, resid loss = {loss_res.item():.3e}, total loss = {loss.item():.3e}")

# ------------------------------------------------------------
# 5. Verify against true zeta (mpmath)
# ------------------------------------------------------------
# Compute true zeta on the critical line using mpmath
def true_zeta_half_it(t_vals):
    return [mpmath.zeta(0.5 + 1j * t) for t in t_vals]

t_test = np.linspace(0.5, 15, 200)
with torch.no_grad():
    zeta_learned = zeta_approx(torch.tensor(t_test)).numpy()
    true_vals = true_zeta_half_it(t_test)
    true_vals = np.array([complex(z) for z in true_vals])

# Plot real part
plt.figure(figsize=(12,5))
plt.subplot(1,2,1)
plt.plot(t_test, zeta_learned.real, label='Learned (wave shaper)', alpha=0.8)
plt.plot(t_test, true_vals.real, '--', label='True ζ(½+it)', alpha=0.8)
plt.xlabel('t (Imaginary part)')
plt.ylabel('Real part')
plt.title('Real part of ζ(½+it)')
plt.legend()

plt.subplot(1,2,2)
plt.plot(t_test, zeta_learned.imag, label='Learned')
plt.plot(t_test, true_vals.imag, '--', label='True')
plt.xlabel('t')
plt.ylabel('Imag part')
plt.title('Imaginary part')
plt.legend()
plt.tight_layout()
plt.show()

# Print error
err = np.mean(np.abs(zeta_learned - true_vals))
print(f"Mean absolute error over t ∈ [0.5,15] : {err:.4e}")