import torch
import torchvision
import torchvision.transforms as transforms

def main():
    # 1. Setup Device and Constants
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    INPUT_DIM = 28 * 28  # MNIST size
    FEATURE_DIM = 4000   # Expanded feature mapping substrate
    NUM_CLASSES = 10
    LAMBDA_REG = 1e-2    # Tikhonov regularization factor (the closure boundary)

    # 2. Load MNIST Dataset
    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)

    # Full-batch processing to compute the global one-shot matrix
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=len(train_dataset), shuffle=False)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=len(test_dataset), shuffle=False)

    x_train, y_train = next(iter(train_loader))
    x_test, y_test = next(iter(test_loader))

    # Flatten the images: (N, 784)
    x_train = x_train.view(-1, INPUT_DIM).to(device)
    x_test = x_test.view(-1, INPUT_DIM).to(device)

    # Convert labels to one-hot encoding targets: (N, 10)
    y_train_onehot = torch.nn.functional.one_hot(y_train, num_classes=NUM_CLASSES).float().to(device)
    y_test = y_test.to(device)

    # 3. Generate the Structural Feature Map Substrate
    # We create a static, random projection matrix to lift the data dimension
    torch.manual_seed(42)
    projection_matrix = torch.randn(INPUT_DIM, FEATURE_DIM, device=device) / (INPUT_DIM ** 0.5)
    
    # Apply non-linear feature mapping
    h_train_raw = torch.relu(torch.matmul(x_train, projection_matrix))
    h_test_raw = torch.relu(torch.matmul(x_test, projection_matrix))

    # 4. Spinorial-Spherical Framework Projection
    # Project features onto the S^{n-1} unit sphere manifold to stabilize the landscape
    h_train = torch.nn.functional.normalize(h_train_raw, p=2, dim=1)
    h_test = torch.nn.functional.normalize(h_test_raw, p=2, dim=1)

    # Add a bias term to the spherical coordinate matrix
    bias_train = torch.ones(h_train.size(0), 1, device=device)
    bias_test = torch.ones(h_test.size(0), 1, device=device)
    X = torch.cat((h_train, bias_train), dim=1)      # Shape: (60000, FEATURE_DIM + 1)
    X_test = torch.cat((h_test, bias_test), dim=1)   # Shape: (10000, FEATURE_DIM + 1)

    print("Computing the analytical one-shot parameter matrix...")
    # 5. One-Shot Analytical Closure Execution
    # Formula: W = (X^T * X + λ * I)^(-1) * X^T * Y
    # This solves the complete optimization manifold instantly.
    XT = X.t()
    XT_X = torch.matmul(XT, X)
    
    # Apply the regularized identity matrix boundary constraint
    identity = torch.eye(X.size(1), device=device)
    restoration_tensor = XT_X + LAMBDA_REG * identity
    
    # Compute the inverse and map directly to the target outputs
    XT_Y = torch.matmul(XT, y_train_onehot)
    
    # W represents our instantly solved One-Shot Parameter Matrix
    W = torch.linalg.solve(restoration_tensor, XT_Y)
    print(f"One-shot Parameter Matrix computed successfully. Shape: {W.shape}")

    # 6. Evaluation (Zero-Training Inference)
    predictions_test = torch.matmul(X_test, W)
    predicted_classes = torch.argmax(predictions_test, dim=1)
    
    accuracy = (predicted_classes == y_test).float().mean().item() * 100
    print(f"\nFinal Evaluation Results:")
    print(f"--------------------------")
    print(f"Test Accuracy: {accuracy:.2f}%")
    print(f"Epochs Required: 0 (Instant mathematical alignment)")

if __name__ == '__main__':
    main()