import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# ─── Define MLP ───────────────────────────────────────────
class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dims, output_dim, dropout=0.3):
        super().__init__()
        layers = []
        dims = [input_dim] + hidden_dims
        for i in range(len(dims) - 1):
            layers += [
                nn.Linear(dims[i], dims[i+1]),
                nn.BatchNorm1d(dims[i+1]),  # normalise activations
                nn.GELU(),                  # smooth non-linearity
                nn.Dropout(dropout),        # regularisation
            ]
        layers.append(nn.Linear(dims[-1], output_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)

# ─── Training Loop ─────────────────────────────────────────
def train(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()        # clear previous gradients
        logits = model(X)            # forward pass
        loss = criterion(logits, y)  # compute loss
        loss.backward()             # backpropagation
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # gradient clip
        optimizer.step()            # update weights
        total_loss += loss.item()
    return total_loss / len(loader)

# ─── Instantiate and run ────────────────────────────────────
device = "cuda" if torch.cuda.is_available() else "cpu"
model = MLP(input_dim=784, hidden_dims=[512, 256, 128], output_dim=10).to(device)
optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
criterion = nn.CrossEntropyLoss()

import torch.nn as nn

class ResidualBlock(nn.Module):
    def __init__(self, channels, stride=1):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, stride=stride, padding=1, bias=False)
        self.bn1   = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2   = nn.BatchNorm2d(channels)
        self.relu  = nn.ReLU(inplace=True)
        # Shortcut if stride changes spatial dims
        self.shortcut = nn.Sequential(
            nn.Conv2d(channels, channels, 1, stride=stride, bias=False),
            nn.BatchNorm2d(channels)
        ) if stride != 1 else nn.Identity()

    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)   # ← skip connection
        return self.relu(out)

# Transfer learning with pretrained ResNet
import torchvision.models as models
backbone = models.resnet50(weights="IMAGENET1K_V1")
backbone.fc = nn.Linear(2048, num_classes)   # replace head

# Freeze backbone, fine-tune head only
for p in backbone.parameters():
    p.requires_grad = False
for p in backbone.fc.parameters():
    p.requires_grad = True

import torch
import torch.nn as nn

class BiLSTMClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            embed_dim, hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,   # forward + backward
            dropout=0.3
        )
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),  # *2 for bidir
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, num_classes)
        )

    def forward(self, x, lengths):
        emb = self.embedding(x)           # (B, T, E)
        # Pack for variable-length sequences
        packed = nn.utils.rnn.pack_padded_sequence(emb, lengths, batch_first=True, enforce_sorted=False)
        out, (hn, _) = self.lstm(packed)
        # Concat last forward + backward hidden states
        last_hidden = torch.cat([hn[-2], hn[-1]], dim=1)  # (B, H*2)
        return self.classifier(last_hidden)

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.h   = num_heads
        self.Wq  = nn.Linear(d_model, d_model)
        self.Wk  = nn.Linear(d_model, d_model)
        self.Wv  = nn.Linear(d_model, d_model)
        self.Wo  = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        B, T, D = q.shape
        Q = self.Wq(q).view(B, T, self.h, self.d_k).transpose(1,2)
        K = self.Wk(k).view(B, -1, self.h, self.d_k).transpose(1,2)
        V = self.Wv(v).view(B, -1, self.h, self.d_k).transpose(1,2)

        # Scaled dot-product attention
        scores = (Q @ K.transpose(-2,-1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = F.softmax(scores, dim=-1)
        out  = (attn @ V).transpose(1,2).reshape(B, T, D)
        return self.Wo(out), attn

import torch.nn as nn

class DCGANGenerator(nn.Module):
    def __init__(self, latent_dim=100, channels=3):
        super().__init__()
        def block(in_c, out_c, stride=2, padding=1):
            return [
                nn.ConvTranspose2d(in_c, out_c, 4, stride, padding, bias=False),
                nn.BatchNorm2d(out_c),
                nn.ReLU(True)
            ]
        self.net = nn.Sequential(
            # 1×1 → 4×4
            nn.ConvTranspose2d(latent_dim, 512, 4, 1, 0, bias=False),
            nn.BatchNorm2d(512), nn.ReLU(True),
            *block(512, 256),  # 8×8
            *block(256, 128),  # 16×16
            *block(128, 64),   # 32×32
            nn.ConvTranspose2d(64, channels, 4, 2, 1, bias=False),
            nn.Tanh()          # 64×64, range [-1,1]
        )
    def forward(self, z):
        return self.net(z.view(-1, z.shape[1], 1, 1))

import torch
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()   # handles FP16 loss scaling

def train_step(model, batch, optimizer, criterion):
    X, y = batch
    optimizer.zero_grad()

    with autocast():                  # FP16 forward pass
        logits = model(X)
        loss   = criterion(logits, y)

    scaler.scale(loss).backward()    # scaled gradients
    scaler.unscale_(optimizer)         # unscale before clip
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)             # update weights
    scaler.update()                    # adjust scale factor
    return loss.item()

# LR warmup + cosine decay (Transformers best practice)
import math

def get_lr(step, d_model, warmup_steps):
    if step == 0: return 0.0
    scale = min(step ** -0.5, step * warmup_steps ** -1.5)
    return d_model ** -0.5 * scale   # original transformer formula

import torch
import onnxruntime as ort
from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np

# ─── Export to ONNX ─────────────────────────────────────────
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, "model.onnx",
    opset_version=17,
    input_names=["image"], output_names=["logits"],
    dynamic_axes={"image": {0: "batch"}}   # variable batch size
)

# ─── ONNX Runtime Inference ─────────────────────────────────
sess_opts = ort.SessionOptions()
sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession("model.onnx", sess_options=sess_opts,
                                providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

# ─── FastAPI endpoint ────────────────────────────────────────
app = FastAPI(title="Deep Learning Model API")

class PredictRequest(BaseModel):
    image: list[list[list[list[float]]]]  # NCHW float array

@app.post("/predict")
async def predict(req: PredictRequest):
    x = np.array(req.image, dtype=np.float32)
    logits = session.run(["logits"], {"image": x})[0]
    probs  = np.exp(logits) / np.exp(logits).sum(-1, keepdims=True)
    top_k  = np.argsort(probs[0])[::-1][:5]
    return {"top5_classes": top_k.tolist(), "probs": probs[0][top_k].tolist()}