import neurograd as ng

# Create tensors with gradient tracking
x = ng.Tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
y = ng.Tensor([[2.0, 1.0], [1.0, 2.0]], requires_grad=True)

# Perform operations
z = x @ y + x.sin()  # Matrix multiplication + element-wise sine
loss = z.sum()       # Scalar loss

# Automatic differentiation
loss.backward()
print(f"x.grad: {x.grad}")

from neurograd.nn.layers.linear import Linear, MLP
from neurograd.nn.losses import MSE
from neurograd.optim.adam import Adam
from neurograd.utils.data import Dataset, DataLoader
from neurograd.nn.metrics import accuracy_score

# Create dataset and model
dataset = Dataset(X_train, y_train)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
model = MLP([784, 128, 64, 10])  # Input -> Hidden -> Output

# Define loss and optimizer
criterion = MSE()
optimizer = Adam(model.named_parameters(), lr=0.001)

# Training loop
for epoch in range(100):
    for X_batch, y_batch in dataloader:
        # Forward pass
        output = model(X_batch)
        loss = criterion(y_batch, output)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # Evaluate
    model.eval()
    pred = model(X_test)
    acc = accuracy_score(y_test, pred)
    model.train()
    print(f"Epoch {epoch}, Loss: {loss.data:.4f}, Accuracy: {acc:.4f}")

from neurograd.amp import autocast, GradScaler
from neurograd.nn.layers.conv import Conv2D, MaxPool2D
from neurograd.nn.layers.linear import Linear
from neurograd.nn.module import Sequential
from neurograd.functions.activations import ReLU, Softmax

# Create CNN model (channels-first: NCHW)
model = Sequential(
    Conv2D(1, 32, kernel_size=3, padding="same", activation="relu"),
    MaxPool2D(pool_size=2),
    Conv2D(32, 64, kernel_size=3, padding="same", activation="relu"),
    MaxPool2D(pool_size=2),
    Flatten(),
    Linear(64 * 7 * 7, 128, activation="relu"),
    Linear(128, 10),
    Softmax(axis=1)
)

# Setup mixed precision
optimizer = Adam(model.named_parameters(), lr=0.001)
loss_fn = CategoricalCrossEntropy()
scaler = GradScaler()

# Training with mixed precision
for epoch in range(num_epochs):
    for batch_x, batch_y in dataloader:
        optimizer.zero_grad()
        
        # Mixed precision forward pass
        with autocast(enabled=True):
            predictions = model(batch_x)        # Auto FP16 where safe
            loss = loss_fn(batch_y, predictions)  # Auto FP32 for stability
        
        # Gradient scaling for FP16 stability
        scaled_loss = scaler.scale(loss)
        scaled_loss.backward()
        scaler.step(optimizer)  # Unscales gradients automatically
        scaler.update()         # Updates scale factor
        
        print(f"Loss: {loss.data.item():.4f}, Scale: {scaler.get_scale():.0f}")

# Benefits: ⚡ 1.5-2x faster, 💾 40-50% less memory, 🎯 same accuracy

# Linear layers with built-in features
layer = Linear(784, 128, activation="relu", dropout=0.2, 
               batch_normalization=True, weights_initializer="he")

# Convolutional layers (channels-first: NCHW)
conv = Conv2D(3, 64, kernel_size=(3,3), padding="same", activation="relu")
pool = MaxPool2D(pool_size=(2,2), strides=(2,2))

# Activations and losses
from neurograd.functions.activations import ReLU, Sigmoid, Softmax
from neurograd.nn.losses import MSE, CategoricalCrossEntropy

# Optimizers
optimizer = Adam(model.named_parameters(), lr=0.001, beta1=0.9, beta2=0.999)
optimizer = SGD(model.named_parameters(), lr=0.01, beta=0.9, weight_decay=1e-4)

x = ng.Tensor([1.0, 2.0, 3.0], requires_grad=True)

# Arithmetic: +, -, *, /, **
z = x + y, x * y, x ** 2

# Math functions
y = x.log(), x.exp(), x.sin(), x.sqrt(), x.abs()

# Linear algebra
C = A @ B           # Matrix multiplication
D = A.transpose()   # Transpose

# Reductions with axis support
s = x.sum(axis=0), x.mean(axis=1, keepdims=True), x.max(), x.std()

from neurograd.utils.data import Dataset, DataLoader

dataset = Dataset(X, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, seed=42)

for batch_idx, (X_batch, y_batch) in enumerate(dataloader):
    output = model(X_batch)
    loss = criterion(y_batch, output)

from neurograd.functions.base import Function

class Swish(Function):
    def forward(self, x):
        self.sigmoid_x = 1 / (1 + ng.xp.exp(-x))
        return x * self.sigmoid_x
    
    def backward(self, grad_output):
        x = self.parent_tensors[0]
        swish_grad = self.sigmoid_x * (1 + x.data * (1 - self.sigmoid_x))
        return grad_output * swish_grad if x.requires_grad else None

from neurograd.utils.grad_check import gradient_check

is_correct = gradient_check(model, X, y, loss_fn, epsilon=1e-7)
print(f"Gradients correct: {is_correct}")

# Visualize computational graphs
fig = loss.visualize_graph(title="Training Loss Graph")
loss.save_graph("computation_graph.png")
loss.print_graph()

# Graph statistics
stats = loss.graph_stats()
print(f"Nodes: {stats['num_tensors']}, Depth: {stats['max_depth']}")

import neurograd as ng

# Save checkpoint
ng.save({
    'model_state': model.state_dict(),
    'optimizer_state': optimizer.state_dict(),
    # optional: 'scaler_state': scaler.state_dict(), 'epoch': epoch
}, 'checkpoint.pth')

# Or use the convenience helper
ng.save_checkpoint(model=model, optimizer=optimizer, path='checkpoint.pth', epoch=epoch)

# Load checkpoint later
ckpt = ng.load('checkpoint.pth')  # or ng.load_checkpoint('checkpoint.pth')
model.load_state_dict(ckpt['model_state'])
optimizer.load_state_dict(ckpt['optimizer_state'])