import numpy as np

class TransferModel:
    """Transfer Learning with frozen backbone"""

    def __init__(self, pretrained_weights, num_classes):
        # Frozen backbone (pre-trained) — DO NOT update!
        self.backbone = pretrained_weights  # dict of W, b per layer
        self.frozen = True

        # New classification head — TRAIN this!
        feat_dim = 512  # output dim of backbone
        self.W_head = np.random.randn(feat_dim, num_classes) * 0.01
        self.b_head = np.zeros((1, num_classes))

    def extract_features(self, X):
        """Run X through frozen backbone"""
        h = X
        for layer in self.backbone:
            h = np.maximum(0, h @ layer['W'] + layer['b'])
        return h  # features (not trained!)

    def forward(self, X):
        features = self.extract_features(X)  # frozen
        logits = features @ self.W_head + self.b_head  # trainable
        exp_l = np.exp(logits - np.max(logits, axis=1, keepdims=True))
        return exp_l / exp_l.sum(axis=1, keepdims=True)

    def train_head(self, X, y_oh, lr=0.01):
        """Only update classification head"""
        features = self.extract_features(X)  # no grad here
        probs = self.forward(X)
        m = X.shape[0]
        dz = (probs - y_oh) / m
        # Only update head weights
        self.W_head -= lr * features.T @ dz
        self.b_head -= lr * np.sum(dz, axis=0, keepdims=True)

# With just 100 labeled samples + pre-trained backbone → 90%+ accuracy!