# nn_torch.py
#
# The same neural network as nn_numpy.py, but using PyTorch.
# Compare this to the numpy version to see what the framework handles for you:
#   - Automatic differentiation (no manual backprop)
#   - Built-in optimizers (Adam instead of hand-coded gradient descent)
#   - GPU support (if available)
#
# CHEG 667-013
# E. M. Furst

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# ── Load and prepare data ──────────────────────────────────────

data = np.loadtxt("data/n2_cp.csv", delimiter=",", skiprows=1)
T_raw = data[:, 0]
Cp_raw = data[:, 1]

# Normalize to [0, 1]
T_min, T_max = T_raw.min(), T_raw.max()
Cp_min, Cp_max = Cp_raw.min(), Cp_raw.max()

X = torch.tensor((T_raw - T_min) / (T_max - T_min), dtype=torch.float32).reshape(-1, 1)
Y = torch.tensor((Cp_raw - Cp_min) / (Cp_max - Cp_min), dtype=torch.float32).reshape(-1, 1)

# ── Define the network ─────────────────────────────────────────
#
# nn.Sequential stacks layers in order. Compare this to nanoGPT's
# model.py, which uses the same PyTorch building blocks (nn.Linear,
# activation functions) but with many more layers.

H = 10  # hidden neurons

model = nn.Sequential(
    nn.Linear(1, H),    # input -> hidden (W1, b1)
    nn.Tanh(),           # activation
    nn.Linear(H, 1),    # hidden -> output (W2, b2)
)

print(f"Model:\n{model}")
print(f"Total parameters: {sum(p.numel() for p in model.parameters())}\n")

# ── Training ───────────────────────────────────────────────────

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()

epochs = 5000
log_interval = 500
losses = []

for epoch in range(epochs):
    # Forward pass -- PyTorch tracks operations for automatic differentiation
    Y_pred = model(X)
    loss = loss_fn(Y_pred, Y)
    losses.append(loss.item())

    # Backward pass -- PyTorch computes all gradients automatically
    optimizer.zero_grad()   # reset gradients from previous step
    loss.backward()         # compute gradients via automatic differentiation
    optimizer.step()        # update weights (Adam optimizer)

    if epoch % log_interval == 0 or epoch == epochs - 1:
        print(f"Epoch {epoch:5d}  Loss: {loss.item():.6f}")

# ── Results ────────────────────────────────────────────────────

# Predict on a fine grid
T_fine = torch.linspace(0, 1, 200).reshape(-1, 1)
with torch.no_grad():  # no gradient tracking needed for inference
    Cp_pred_norm = model(T_fine)

# Convert back to physical units
T_fine_K = T_fine.numpy() * (T_max - T_min) + T_min
Cp_pred = Cp_pred_norm.numpy() * (Cp_max - Cp_min) + Cp_min

# ── Plot ───────────────────────────────────────────────────────

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

ax1.plot(T_raw, Cp_raw, 'ko', markersize=6, label='NIST data')
ax1.plot(T_fine_K, Cp_pred, 'r-', linewidth=2, label=f'NN ({H} neurons)')
ax1.set_xlabel('Temperature (K)')
ax1.set_ylabel('$C_p$ (kJ/kg/K)')
ax1.set_title('$C_p(T)$ for N$_2$ at 1 bar — PyTorch')
ax1.legend()

ax2.semilogy(losses)
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Mean Squared Error')
ax2.set_title('Training Loss')

plt.tight_layout()
plt.savefig('nn_fit_torch.png', dpi=150)
plt.show()