import torch

print(torch.__version__)                        # PyTorch version string
print(torch.version.cuda)                       # Compiled CUDA version
print(torch.cuda.is_available())                # True if NVIDIA GPU with CUDA is ready
print(torch.cuda.device_count())                # Number of accessible GPUs
print(torch.cuda.get_device_name(0))            # GPU name (e.g., 'NVIDIA A100-SXM4-80GB')
print(torch.backends.mps.is_available())        # True if Apple Silicon GPU (Metal) is available

# Automatic device selection: CUDA -> Apple Silicon (MPS) -> CPU fallback
device = torch.device(
    'cuda' if torch.cuda.is_available()
    else 'mps' if torch.backends.mps.is_available()
    else 'cpu'
)
print(f"Using device: {device}")

# From Python lists
x = torch.tensor([1, 2, 3])                     # 1D tensor of int64
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])      # 2D tensor of float32

# From / to NumPy arrays (shared memory buffer)
import numpy as np
arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
t = torch.from_numpy(arr)                       # Zero-copy tensor from NumPy
np_arr = t.numpy()                              # Back to NumPy (CPU tensor only)

x = torch.zeros(2, 3)                           # 2x3 tensor filled with 0.0
x = torch.ones(2, 3)                            # 2x3 tensor filled with 1.0
x = torch.empty(2, 3)                           # Uninitialized memory allocation
x = torch.full((2, 3), 3.14159)                 # Filled with specified scalar value
x = torch.eye(4)                                # 4x4 Identity matrix (diagonal 1s)

# Like-operators: Match shape, dtype, and device of reference tensor
y = torch.zeros_like(x)                         # Same shape and type as x, filled with 0
y = torch.ones_like(x)                          # Same shape and type as x, filled with 1
y = torch.empty_like(x)                         # Same shape and type as x, uninitialized

x = torch.arange(0, 10, step=2)                 # [0, 2, 4, 6, 8] (half-open [start, end))
x = torch.linspace(start=0, end=1, steps=5)     # [0.0, 0.25, 0.5, 0.75, 1.0] (inclusive)
x = torch.logspace(start=0, end=3, steps=4)     # [1., 10., 100., 1000.] (base 10)

torch.manual_seed(42)                           # Set CPU random seed for reproducibility
torch.cuda.manual_seed_all(42)                  # Set random seed across all GPUs

x = torch.rand(2, 3)                            # Uniform distribution U[0, 1)
x = torch.randn(2, 3)                           # Standard normal distribution N(0, 1)
x = torch.randint(low=0, high=10, size=(2, 3))  # Discrete uniform integers in [low, high)
x = torch.randperm(10)                          # Random permutation of integers 0 to 9

t = torch.randn(3, 4, 5, dtype=torch.float32, device='cpu')

print(t.shape)                                  # torch.Size([3, 4, 5])
print(t.size())                                 # Equivalent to t.shape
print(t.ndim)                                   # Number of dimensions: 3
print(t.numel())                                # Total elements: 3 * 4 * 5 = 60
print(t.dtype)                                  # torch.float32
print(t.device)                                 # cpu
print(t.requires_grad)                          # False (autograd tracking flag)
print(t.is_cuda)                                # False

x = torch.tensor([1, 2, 3])

# Explicit casting via .to()
x_float = x.to(torch.float32)                   # Cast to float32 (single precision)
x_half  = x.to(torch.float16)                   # Cast to float16 (half precision)
x_bf16  = x.to(torch.bfloat16)                  # Cast to brain floating point
x_int   = x.to(torch.int32)                     # Cast to int32
x_long  = x.to(torch.int64)                     # Cast to int64 (indices / token IDs)

# Convenience type casting methods
x_f = x.float()                                 # Shorthand for .to(torch.float32)
x_d = x.double()                                # Shorthand for .to(torch.float64)
x_l = x.long()                                  # Shorthand for .to(torch.int64)
x_b = x.bool()                                  # Shorthand for .to(torch.bool)

x = torch.randn(2, 3)

x_gpu = x.to(device)                            # Move to active compute device
x_cpu = x_gpu.to('cpu')                         # Move back to host system memory
x_gpu = x.cuda(0)                               # Move specifically to GPU 0
x_cpu = x_gpu.cpu()                             # Shorthand for .to('cpu')

x = torch.arange(24).reshape(2, 3, 4)

# Multi-dimensional slicing
sub = x[0, :, :]                                # First batch item: shape [3, 4]
sub = x[:, 1, :]                                # Row 1 across all batches: shape [2, 4]
sub = x[..., -1]                                # Last column across all leading dims: shape [2, 3]

# Boolean masking & filtering
mask = x > 12                                   # Boolean condition tensor
filtered = x[mask]                              # 1D tensor with values satisfying condition
replaced = torch.where(x > 10, x, torch.zeros_like(x)) # Conditional replacement: if cond then x else 0

x = torch.randn(2, 3, 4)

# Reshaping
y = x.view(2, 12)                               # Reshape (requires contiguous memory layout)
y = x.reshape(2, 12)                            # Safe reshape (clones if non-contiguous)
y = x.flatten(start_dim=1)                      # Flatten to [2, 12] (standard for CNN heads)

# Squeeze & Unsqueeze (dimension expansion/reduction)
x_sq   = x.unsqueeze(dim=0)                     # Shape: [1, 2, 3, 4] (add batch dimension)
x_sq2  = x[:, None, :, :]                       # Equivalent unsqueeze via None indexing
x_red  = x_sq.squeeze(dim=0)                    # Shape: [2, 3, 4] (removes dimensions of size 1)

# Permute & Transpose
y = x.transpose(1, 2)                           # Swaps dim 1 and dim 2: shape [2, 4, 3]
y = x.permute(2, 0, 1)                          # Arbitrary axis reordering: shape [4, 2, 3]
y_cont = y.contiguous()                         # Make memory contiguous after permute

a = torch.randn(2, 3)
b = torch.randn(2, 3)

# Concatenate along existing dimension
c = torch.cat([a, b], dim=0)                    # Shape: [4, 3] (vertical concat)
c = torch.cat([a, b], dim=1)                    # Shape: [2, 6] (horizontal concat)

# Stack along a new dimension
s = torch.stack([a, b], dim=0)                  # Shape: [2, 2, 3] (adds new dim)

# Splitting & Chunking
chunks = torch.chunk(c, chunks=2, dim=0)        # Split into 2 equal chunks
splits = torch.split(c, [1, 3], dim=0)          # Split into custom section sizes [1, 3]

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])

# Out-of-place arithmetic
c = a + b                                       # Element-wise addition
c = a * b                                       # Element-wise multiplication (Hadamard product)
c = a / b                                       # Element-wise division
c = torch.pow(a, 2)                             # Exponentiation (a ** 2)
c = torch.exp(a)                                # Natural exponential e^a
c = torch.log(a)                                # Natural logarithm ln(a)
c = torch.sqrt(a)                               # Square root
c = torch.clamp(a, min=1.5, max=2.5)            # Clip values between min and max

# In-place operations (trailing underscore modifies tensor directly)
a.add_(b)                                       # a = a + b in-place (saves memory allocations)
a.mul_(2.0)                                     # a = a * 2.0 in-place
a.zero_()                                       # Fill a with zeros in-place

x = torch.randn(4, 5)

total   = torch.sum(x)                          # Scalar sum of all elements
row_sum = torch.sum(x, dim=1, keepdim=True)     # Sum along columns: shape [4, 1]
mean    = torch.mean(x, dim=0)                  # Mean along rows: shape [5]
std     = torch.std(x)                          # Standard deviation
max_val, max_idx = torch.max(x, dim=1)          # Values and argmax indices along dim 1
argmax  = torch.argmax(x, dim=-1)               # Best token index across vocabulary logits

# 1. 1D Vector Dot Product
u = torch.randn(4)
v = torch.randn(4)
dot = torch.dot(u, v)                           # Scalar dot product: u^T v

# 2. 2D Matrix Multiplication
A = torch.randn(3, 4)
B = torch.randn(4, 5)
C = torch.mm(A, B)                              # Matrix product: shape [3, 5]
C = A @ B                                       # Python matrix operator syntax

# 3. Batched Matrix Multiplication (BMM)
B_A = torch.randn(16, 3, 4)                     # Batch size 16
B_B = torch.randn(16, 4, 5)
B_C = torch.bmm(B_A, B_B)                       # Batched matmul: shape [16, 3, 5]
B_C = B_A @ B_B                                 # Supports arbitrary leading batch dimensions

# 4. Einstein Summation (einsum)
# Multi-head attention computation: Q @ K^T
# Q: [batch, heads, seq_q, d_k], K: [batch, heads, seq_k, d_k]
q = torch.randn(2, 8, 64, 32)
k = torch.randn(2, 8, 64, 32)
scores = torch.einsum('b h i d, b h j d -> b h i j', q, k) # Shape [2, 8, 64, 64]

x = torch.tensor([2.0, 3.0], requires_grad=True)

# Forward pass
y = x ** 2 + 3 * x + 1
out = y.sum()                                   # Scalar output: (2^2 + 6 + 1) + (3^2 + 9 + 1) = 11 + 19 = 30

# Backward pass (compute gradients d(out)/dx)
out.backward()

print(x.grad)                                   # dy/dx = 2x + 3 -> [7.0, 9.0]

# Context 1: torch.no_grad()
# Disables graph creation to save memory during evaluation / validation
with torch.no_grad():
    predictions = model(test_inputs)

# Context 2: torch.inference_mode()
# Even faster than no_grad(); disables version tracking and view tracking
with torch.inference_mode():
    features = model(test_inputs)

# Detach tensor from computation graph
detached_tensor = y.detach()                    # Shares data, but stops gradient propagation

# Gradients accumulate by default in PyTorch!
# Must zero out gradient buffers before each backprop step:
optimizer.zero_grad(set_to_none=True)           # set_to_none=True reduces memory bandwidth

# Retain gradient on intermediate activations (non-leaf tensors)
h = x * 3
h.retain_grad()                                 # Allows inspecting h.grad after backward()

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiLayerPerceptron(nn.Module):
    def __init__(self, in_features: int, hidden_dim: int, num_classes: int, dropout_p: float = 0.2):
        super().__init__()
        # Parameterized sub-modules
        self.fc1 = nn.Linear(in_features, hidden_dim)
        self.norm = nn.LayerNorm(hidden_dim)
        self.dropout = nn.Dropout(dropout_p)
        self.fc2 = nn.Linear(hidden_dim, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Forward computational path
        x = self.fc1(x)
        x = self.norm(x)
        x = F.gelu(x)                            # Functional activation
        x = self.dropout(x)
        logits = self.fc2(x)
        return logits

# Instantiate and verify parameter counts
model = MultiLayerPerceptron(in_features=128, hidden_dim=256, num_classes=10)
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable parameters: {total_params:,}")

# Linear / Dense
linear = nn.Linear(in_features=768, out_features=768, bias=True)

# Convolutions
conv2d = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
deconv = nn.ConvTranspose2d(in_channels=64, out_channels=3, kernel_size=4, stride=2, padding=1)

# Normalization
layer_norm = nn.LayerNorm(normalized_shape=768)
batch_norm = nn.BatchNorm2d(num_features=64)
rms_norm   = nn.RMSNorm(normalized_shape=768)    # PyTorch 2.4+ standard for LLMs

# Pooling
max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
avg_pool = nn.AdaptiveAvgPool2d(output_size=(1, 1)) # Downsample to exactly 1x1 spatial

# NLP & Embeddings
embedding = nn.Embedding(num_embeddings=50000, embedding_dim=768)
mha = nn.MultiheadAttention(embed_dim=768, num_heads=12, batch_first=True)

# Container modules
seq_model = nn.Sequential(
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)

layer_list = nn.ModuleList([nn.Linear(64, 64) for _ in range(4)])
layer_dict = nn.ModuleDict({
    'conv': nn.Conv2d(3, 16, 3),
    'fc': nn.Linear(16, 10),
})

model.train()                                   # Activates Dropout & BatchNorm training updates
model.eval()                                    # Disables Dropout & locks BatchNorm running stats

# Classification: Multi-class (logits -> labels)
# Note: CrossEntropyLoss applies LogSoftmax internally; input raw logits!
criterion_ce = nn.CrossEntropyLoss(label_smoothing=0.1)
loss = criterion_ce(logits, target_class_indices) # logits: [B, C], targets: [B]

# Classification: Binary / Multi-label
criterion_bce = nn.BCEWithLogitsLoss()          # Numerically stable Sigmoid + BCE
loss = criterion_bce(binary_logits, binary_targets)

# Regression: Mean Squared Error (MSE) & L1
criterion_mse = nn.MSELoss()
criterion_l1  = nn.L1Loss()
criterion_hub = nn.SmoothL1Loss(beta=1.0)       # Huber loss: robust against outliers

import torch.optim as optim

# 1. AdamW (Decoupled weight decay — recommended for Transformers & Modern DNNs)
optimizer = optim.AdamW(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,                          # True decoupled L2 regularization
    fused=True if torch.cuda.is_available() else False, # Fused GPU CUDA kernel
)

# 2. Stochastic Gradient Descent with Nesterov Momentum
optimizer_sgd = optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,
    nesterov=True,
    weight_decay=1e-4,
)

# 3. Parameter Groups (Fine-tuning with differential learning rates)
optimizer_groups = optim.AdamW([
    {'params': model.fc1.parameters(), 'lr': 1e-5},
    {'params': model.fc2.parameters(), 'lr': 1e-3, 'weight_decay': 0.0},
])

# Cosine Annealing with warmup
scheduler = optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=100,                                  # Total training epochs
    eta_min=1e-6,                               # Minimum learning rate target
)

# Adaptive decay on metric plateau
scheduler_plateau = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=0.5,                                 # lr = lr * 0.5
    patience=5,                                 # Number of non-improving epochs
)

# Step scheduler after each training epoch:
# scheduler.step()
# scheduler_plateau.step(val_loss)

from torch.utils.data import Dataset, DataLoader, TensorDataset, random_split
import torch

# 1. Quick TensorDataset (for tensors already in memory)
x_data = torch.randn(1000, 32)
y_data = torch.randint(0, 2, (1000,))
dataset = TensorDataset(x_data, y_data)

# Split into train/validation sets (80% / 20%)
train_set, val_set = random_split(dataset, [800, 200])

# 2. Custom Dataset Blueprint
class CustomDataset(Dataset):
    def __init__(self, file_paths, labels, transform=None):
        self.file_paths = file_paths
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.file_paths)

    def __getitem__(self, idx: int):
        # Load sample lazily from disk
        sample = torch.load(self.file_paths[idx])
        label = self.labels[idx]
        if self.transform:
            sample = self.transform(sample)
        return sample, label

# 3. High-Performance DataLoader
loader = DataLoader(
    train_set,
    batch_size=64,
    shuffle=True,                               # Shuffle data every epoch
    num_workers=4,                              # Multiprocess worker subprocesses
    pin_memory=True,                            # Page-locked RAM memory for faster GPU copies
    drop_last=True,                             # Discard incomplete trailing batch
)

def train_one_epoch(model, loader, criterion, optimizer, device, max_grad_norm=1.0):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_idx, (inputs, targets) in enumerate(loader):
        inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)

        # 1. Reset gradients
        optimizer.zero_grad(set_to_none=True)

        # 2. Forward pass
        outputs = model(inputs)
        loss = criterion(outputs, targets)

        # 3. Backward propagation
        loss.backward()

        # 4. Gradient Clipping (prevents exploding gradients)
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)

        # 5. Optimizer parameter update
        optimizer.step()

        # Tracking metrics
        running_loss += loss.item() * inputs.size(0)
        _, preds = torch.max(outputs, 1)
        correct += (preds == targets).sum().item()
        total += targets.size(0)

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc


def evaluate(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.inference_mode():
        for inputs, targets in loader:
            inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            running_loss += loss.item() * inputs.size(0)
            _, preds = torch.max(outputs, 1)
            correct += (preds == targets).sum().item()
            total += targets.size(0)

    return running_loss / total, correct / total

# 1. Recommended: Save only learned parameters (state_dict)
torch.save(model.state_dict(), 'model_weights.pth')

# Restore weights into initialized architecture
model = MultiLayerPerceptron(in_features=128, hidden_dim=256, num_classes=10)
model.load_state_dict(torch.load('model_weights.pth', map_location=device))

# 2. Full Training Checkpoint (resumable training state)
checkpoint = {
    'epoch': 42,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'best_val_loss': 0.1245,
}
torch.save(checkpoint, 'checkpoint_epoch_42.pt')

# Resume checkpoint
loaded = torch.load('checkpoint_epoch_42.pt', map_location=device)
model.load_state_dict(loaded['model_state_dict'])
optimizer.load_state_dict(loaded['optimizer_state_dict'])
scheduler.load_state_dict(loaded['scheduler_state_dict'])
start_epoch = loaded['epoch'] + 1

from torch.amp import autocast, GradScaler

scaler = GradScaler('cuda')                     # Dynamic loss scaler to prevent underflow

for inputs, targets in loader:
    inputs, targets = inputs.to('cuda'), targets.to('cuda')
    optimizer.zero_grad(set_to_none=True)

    # Automatic mixed-precision forward pass
    with autocast('cuda', dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # Scaled backward pass
    scaler.scale(loss).backward()

    # Unscale before clipping (if clipping gradients)
    scaler.unscale_(optimizer)
    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    # Step optimizer and update scaler factor
    scaler.step(optimizer)
    scaler.update()

# JIT-compile model graph with TorchDynamo & Inductor
compiled_model = torch.compile(
    model,
    mode='default',                             # Options: 'default', 'reduce-overhead', 'max-autotune'
)

# Inspect GPU memory allocation
allocated = torch.cuda.memory_allocated() / (1024 ** 2)  # In Megabytes
reserved  = torch.cuda.memory_reserved()  / (1024 ** 2)  # In Megabytes
print(f"Allocated: {allocated:.1f} MB | Reserved: {reserved:.1f} MB")

# Free cached unused memory blocks back to CUDA driver
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()