import torch
import torch.optim as optim
import matplotlib.pyplot as plt

# ------------------------------------------------------------
# 1. Random complex Fourier basis (frozen wave shaper)
# ------------------------------------------------------------
def random_complex_basis(t, K, omega_scale):
    """
    t: (N,) tensor of collocation points
    Returns: (N, K) complex tensor, each column = exp(i * omega_k * t)
    """
    omega = torch.randn(K) * omega_scale
    # shape (K, N) -> transpose to (N, K)
    phi = torch.exp(1j * torch.outer(t, omega))
    return phi

# ------------------------------------------------------------
# 2. ODE: Duffing oscillator u'' + delta*u' + alpha*u + beta*u^3 = 0
# ------------------------------------------------------------
def duffing_residual(u, t, delta=0.2, alpha=-1.0, beta=1.0):
    # u is (N,) real
    u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]
    u_tt = torch.autograd.grad(u_t, t, grad_outputs=torch.ones_like(u_t), create_graph=True)[0]
    return u_tt + delta*u_t + alpha*u + beta*u**3

# ------------------------------------------------------------
# 3. Set up problem
# ------------------------------------------------------------
t_colloc = torch.linspace(0, 10, 200, requires_grad=True)  # collocation points
K = 100                    # number of basis functions (large)
omega_scale = 5.0          # random frequency range

phi_colloc = random_complex_basis(t_colloc, K, omega_scale)  # (200, K) complex

# Initial conditions: u(0)=1, u'(0)=0
# Build a particular solution u_p that satisfies ICs:
#   u_p(t) = A cos(ω0 t) ? Instead, we solve for coefficients with constraints.
# Simpler: enforce ICs via two extra equations in loss.

# Coefficients to learn (complex)
c = torch.nn.Parameter(torch.randn(K, dtype=torch.complex64) * 0.1)

def u_pred(t):
    phi_t = random_complex_basis(t, K, omega_scale)  # rebuild for new t
    return torch.real(phi_t @ c)  # take real part for physical solution

# ------------------------------------------------------------
# 4. Loss: ODE residual + initial conditions
# ------------------------------------------------------------
def loss_fn():
    u_colloc = torch.real(phi_colloc @ c)  # (200,) real
    # Need gradients through t_colloc
    # But phi_colloc is frozen; u_colloc is real part of linear combination
    # To compute derivatives, we must recompute phi with autograd.
    # Simpler: use t_colloc with requires_grad=True and recompute u(t)
    u_t = torch.real(phi_colloc @ c)  # but this breaks gradient chain because phi_colloc is detached?
    # Better: define a function of t using the same random frequencies
    # We'll rebuild phi inside a forward pass to keep autograd.
    def u_of_t(t):
        omega = phi_colloc_omega  # we need to store omega
        # For brevity, I'll show a clean version below.
        pass

# Instead of inline complications, here's a minimal working version
# using a pre-defined omega vector and re-computing phi with autograd.

omega_fixed = torch.randn(K) * omega_scale  # fixed random frequencies
omega_fixed.requires_grad_(False)

def u_with_grad(t):
    # t: tensor with requires_grad=True
    phi = torch.exp(1j * torch.outer(t, omega_fixed))  # (len(t), K) complex
    return torch.real(phi @ c)

# Loss
def compute_loss():
    u_colloc = u_with_grad(t_colloc)
    residual = duffing_residual(u_colloc, t_colloc)
    # Initial conditions at t=0
    t0 = torch.tensor([0.0], requires_grad=True)
    u0 = u_with_grad(t0)
    u0_t = torch.autograd.grad(u0, t0, create_graph=True)[0]
    loss_ic = (u0 - 1.0)**2 + (u0_t - 0.0)**2
    loss_res = torch.mean(residual**2)
    return loss_res + 100.0 * loss_ic  # weight ICs

# ------------------------------------------------------------
# 5. Optimize coefficients (snap to solution)
# ------------------------------------------------------------
optimizer = optim.Adam([c], lr=0.01)
for epoch in range(2000):
    optimizer.zero_grad()
    loss = compute_loss()
    loss.backward()
    optimizer.step()
    if epoch % 500 == 0:
        print(f"Epoch {epoch}, loss = {loss.item():.3e}")

# ------------------------------------------------------------
# 6. Plot the learned solution
# ------------------------------------------------------------
t_test = torch.linspace(0, 10, 500)
with torch.no_grad():
    u_test = u_with_grad(t_test).numpy()
plt.plot(t_test.numpy(), u_test)
plt.xlabel('t'); plt.ylabel('u(t)'); plt.title('Duffing oscillator – complex wave solution')
plt.show()