"""Reference run for the guide "A standing wave" (1D wave equation).

u_tt = c^2 u_xx on x in [0,1], t in [0,T], c = 1
u(x,0) = sin(pi x), u_t(x,0) = 0, u(0,t) = u(1,t) = 0
Exact: u = cos(c pi t) sin(pi x)

Hand-written for checking the guide; not the Studio's generated output.
Usage: python wave_1d_reference.py [T] [epochs] [w_ic] [w_bc]
"""
import sys, math
import numpy as np, torch, torch.nn as nn

C = 1.0
T = float(sys.argv[1]) if len(sys.argv) > 1 else 1.0
epochs = int(sys.argv[2]) if len(sys.argv) > 2 else 6000
w_ic = float(sys.argv[3]) if len(sys.argv) > 3 else 10.0
w_bc = float(sys.argv[4]) if len(sys.argv) > 4 else 10.0
tag = sys.argv[5] if len(sys.argv) > 5 else f"T{T:g}"
torch.manual_seed(0); np.random.seed(0)

def exact(x, t): return np.cos(C * math.pi * t) * np.sin(math.pi * x)

net = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(),
                    nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
for m in net:
    if isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.weight); nn.init.zeros_(m.bias)

Nr, Nb, Ni = 3000, 100, 100
xt = torch.rand(Nr, 2); xt[:, 1] *= T
xi = torch.rand(Ni, 1)
tb = torch.rand(Nb, 1) * T

def u_of(x, t): return net(torch.cat([x, t], 1))

def losses():
    x = xt[:, :1].clone().requires_grad_(True); t = xt[:, 1:].clone().requires_grad_(True)
    u = u_of(x, t); one = torch.ones_like(u)
    u_t = torch.autograd.grad(u, t, one, create_graph=True)[0]
    u_tt = torch.autograd.grad(u_t, t, one, create_graph=True)[0]
    u_x = torch.autograd.grad(u, x, one, create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, one, create_graph=True)[0]
    l_pde = ((u_tt - C**2 * u_xx) ** 2).mean()
    # initial displacement and initial velocity at t = 0
    x0 = xi.clone(); t0 = torch.zeros_like(xi).requires_grad_(True)
    u0 = u_of(x0, t0)
    v0 = torch.autograd.grad(u0, t0, torch.ones_like(u0), create_graph=True)[0]
    l_u0 = ((u0 - torch.sin(math.pi * xi)) ** 2).mean()
    l_v0 = (v0 ** 2).mean()
    l_bc = (u_of(torch.zeros_like(tb), tb) ** 2).mean() + (u_of(torch.ones_like(tb), tb) ** 2).mean()
    return l_pde, l_u0, l_v0, l_bc

ng = 101
xg, tg = np.meshgrid(np.linspace(0, 1, ng), np.linspace(0, T, ng))
XT = torch.tensor(np.stack([xg.ravel(), tg.ravel()], 1), dtype=torch.float32)
UE = exact(xg.ravel(), tg.ravel())
def rel_err():
    with torch.no_grad(): up = net(XT).numpy().ravel()
    return float(np.linalg.norm(up - UE) / np.linalg.norm(UE)), up

opt = torch.optim.Adam(net.parameters(), lr=1e-3)
hist = []
for ep in range(epochs + 1):
    opt.zero_grad()
    lp, lu, lv, lb = losses()
    loss = lp + w_ic * (lu + lv) + w_bc * lb
    if ep % 100 == 0 or ep == epochs:
        e, up = rel_err()
        hist.append((ep, loss.item(), lp.item(), lu.item(), lv.item(), lb.item(), e))
        if ep % 500 == 0:
            print(f"ep {ep:5d} loss {loss.item():.3e} pde {lp.item():.2e} u0 {lu.item():.2e} v0 {lv.item():.2e} bc {lb.item():.2e} relL2 {e:.3e}", flush=True)
    if ep < epochs:
        loss.backward(); opt.step()
e, up = rel_err()
print("final relL2", e, "max abs err", np.abs(up - UE).max())
np.savez(f"wave_{tag}.npz", hist=np.array(hist), up=up, ue=UE, xg=xg, tg=tg)