import torch
import torchvision
import torchvision.transforms as transforms


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}\n")

    INPUT_DIM = 28 * 28
    FEATURE_DIM = 1500
    NUM_CLASSES = 10
    LAMBDA_REG = 1e-2
    EPOCHS = 10

    # Phase 1 uses a large batch (it's just accumulating XtX / XtY).
    # Phase 2 is real SGD, so it wants a smaller, shuffled batch.
    BASE_BATCH = 1500
    TRAIN_BATCH = 256
    PULL_LR = 0.5          # was 2e-3 -- far too small for unit-norm features
    WEIGHT_DECAY = 1e-4    # keeps the correction from drifting

    # 1. Load MNIST
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    train_dataset = torchvision.datasets.MNIST(root='../data', train=True, download=True, transform=transform)
    test_dataset = torchvision.datasets.MNIST(root='../data', train=False, download=True, transform=transform)

    base_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BASE_BATCH, shuffle=False)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=TRAIN_BATCH, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BASE_BATCH, shuffle=False)

    # 2. Static random feature map
    torch.manual_seed(42)
    projection_matrix = torch.randn(INPUT_DIM, FEATURE_DIM, device=device) / (INPUT_DIM ** 0.5)

    def make_features(x_flat):
        h = torch.nn.functional.normalize(torch.relu(torch.matmul(x_flat, projection_matrix)), p=2, dim=1)
        return torch.cat((h, torch.ones(h.size(0), 1, device=device)), dim=1)  # append bias column

    # 3. Phase 1: closed-form ridge solution -> the locked base classifier
    XT_X_accum = torch.zeros(FEATURE_DIM + 1, FEATURE_DIM + 1, device=device)
    XT_Y_accum = torch.zeros(FEATURE_DIM + 1, NUM_CLASSES, device=device)

    print("Phase 1: Locking Base Structural Matrix...")
    for x_batch, y_batch in base_loader:
        x_batch = x_batch.view(-1, INPUT_DIM).to(device)
        y_onehot = torch.nn.functional.one_hot(y_batch, NUM_CLASSES).float().to(device)
        X_batch = make_features(x_batch)
        XT_X_accum.add_(torch.matmul(X_batch.t(), X_batch))
        XT_Y_accum.add_(torch.matmul(X_batch.t(), y_onehot))

    identity = torch.eye(FEATURE_DIM + 1, device=device)
    W_base = torch.linalg.solve(XT_X_accum + LAMBDA_REG * identity, XT_Y_accum)
    print("-> Base One-Shot Matrix locked.\n")

    # 4. Phase 2: additive logit-space corrector, W_base stays frozen.
    #    Effective classifier = W_base + W_delta. We learn W_delta only.
    W_delta = torch.zeros(FEATURE_DIM + 1, NUM_CLASSES, device=device)

    print(f"Phase 2: Training the logit-space corrector over {EPOCHS} epochs...")
    print("-" * 75)
    print(f"{'Epoch':<8} | {'Base Accuracy':<18} | {'Corrected Accuracy':<22}")
    print("-" * 75)

    for epoch in range(1, EPOCHS + 1):
        # --- Train: manual cross-entropy gradient on W_delta ---
        with torch.no_grad():
            for x_batch, y_batch in train_loader:
                x_batch = x_batch.view(-1, INPUT_DIM).to(device)
                y_onehot = torch.nn.functional.one_hot(y_batch, NUM_CLASSES).float().to(device)
                X_batch = make_features(x_batch)

                logits = torch.matmul(X_batch, W_base + W_delta)
                probs = torch.nn.functional.softmax(logits, dim=1)

                # dL/d(logits) = (probs - y) for softmax cross-entropy
                grad = torch.matmul(X_batch.t(), probs - y_onehot) / X_batch.size(0)
                grad.add_(WEIGHT_DECAY * W_delta)

                W_delta.sub_(PULL_LR * grad)

        # --- Evaluate base vs corrected ---
        total = correct_base = correct_corr = 0
        with torch.no_grad():
            for x_test, y_test in test_loader:
                x_test = x_test.view(-1, INPUT_DIM).to(device)
                y_test = y_test.to(device)
                X_test = make_features(x_test)

                correct_base += (torch.matmul(X_test, W_base).argmax(1) == y_test).sum().item()
                correct_corr += (torch.matmul(X_test, W_base + W_delta).argmax(1) == y_test).sum().item()
                total += y_test.size(0)

        acc_base = correct_base / total * 100
        acc_corr = correct_corr / total * 100
        print(f"Epoch {epoch:02d}   | {acc_base:.2f}%            | {acc_corr:.2f}%")

    print("-" * 75)
    print("The base row is constant (W_base is fixed by the closed-form solve).")
    print("The corrected row moves because W_delta does real cross-entropy descent")
    print("on the same frozen features. Gains are bounded by how much cross-entropy")
    print("disagrees with the ridge MSE solution -- it is still a linear classifier.")


if __name__ == '__main__':
    main()