import torch
import torch.optim as optim

# Fixed random frequencies (wave shaper)
K = 500
omega = torch.randn(K) * 10.0   # random frequencies, range ~ [-30,30]

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

def zeta_approx(t):
    # t: tensor of real values (imaginary part)
    # returns ζ(½+it) as complex tensor
    phase = torch.outer(t, omega)          # (len(t), K)
    phi = torch.exp(1j * phase)            # basis functions
    return phi @ c                          # linear combination

# Collocation points
t_colloc = torch.linspace(0, 30, 500)

# Loss: functional equation residual
def functional_residual():
    s = 0.5 + 1j * t_colloc
    zeta_s = zeta_approx(t_colloc)
    # compute ζ(1-s) = ζ(0.5 - i t)
    zeta_1_minus_s = zeta_approx(-t_colloc)   # careful: need conjugate? Actually ζ(conj(s)) = conj(ζ(s))
    # Right-hand side of functional equation
    rhs = (2**s) * (torch.pi**(s-1)) * torch.sin(torch.pi * s / 2) * torch.lgamma(1-s).exp() * zeta_1_minus_s
    return torch.mean(torch.abs(zeta_s - rhs)**2)

# Additional constraints: known values
t0 = torch.tensor([0.0])
zeta_half = zeta_approx(t0)   # should be ≈ -1.46035 (real)
loss_value = torch.abs(zeta_half + 1.46035)**2

# Optimize
optimizer = optim.Adam([c], lr=0.01)
for step in range(5000):
    optimizer.zero_grad()
    loss = functional_residual() + 10.0 * loss_value
    loss.backward()
    optimizer.step()