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}")
    
    INPUT_DIM = 28 * 28  
    FEATURE_DIM = 100   # Can safely scale this higher now due to memory optimization
    NUM_CLASSES = 10
    LAMBDA_REG = 1e-2    
    BATCH_SIZE = 100    # 4000 Process chunks at a time to minimize peak memory usage

    # 1. 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)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

    # 2. Static Feature Map Substrate Projection Matrix
    torch.manual_seed(42)
    projection_matrix = torch.randn(INPUT_DIM, FEATURE_DIM, device=device) / (INPUT_DIM ** 0.5)

    # Allocating fixed-size accumulation tensors for the final analytical closure
    # Total matrix size to invert is always exactly (FEATURE_DIM + 1, FEATURE_DIM + 1)
    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("Streaming training data to build analytical landscape incrementally...")

    for _ in range(2):
        # 3. Stream and accumulate the dataset parameters (Memory Bottleneck Solution)
        for x_batch, y_batch in train_loader:
            x_batch = x_batch.view(-1, INPUT_DIM).to(device)
            y_batch_onehot = torch.nn.functional.one_hot(y_batch, num_classes=NUM_CLASSES).float().to(device)
            
            # Project and apply non-linear mapping
            h_batch = torch.relu(torch.matmul(x_batch, projection_matrix))
            
            # Spinorial-Spherical Framework Projection (Unit-Sphere Constraint)
            h_batch = torch.nn.functional.normalize(h_batch, p=2, dim=1)
            
            # Append Bias
            bias = torch.ones(h_batch.size(0), 1, device=device)
            X_batch = torch.cat((h_batch, bias), dim=1)
            
            # Incrementally add to global mathematical system matrices
            XT_X_accum.add_ (torch.matmul(X_batch.t(), X_batch))
            XT_Y_accum.add_ (torch.matmul(X_batch.t(), y_batch_onehot))

    print("Executing One-Shot Analytical Closure...")
    # 4. Global Minimum-Norm Collapse
    identity = torch.eye(FEATURE_DIM + 1, device=device)
    restoration_tensor = XT_X_accum + LAMBDA_REG * identity
    
    # Instant Parameter Matrix Generation
    W = torch.linalg.solve(restoration_tensor, XT_Y_accum)
    print(f"One-shot Parameter Matrix computed. Matrix Shape: {W.shape}")

    # 5. Evaluate incrementally to keep Test Memory ultra-low
    print("\nRunning zero-training inference evaluation...")
    correct = 0
    total = 0
    
    with torch.no_grad():
        for x_test_batch, y_test_batch in test_loader:
            x_test_batch = x_test_batch.view(-1, INPUT_DIM).to(device)
            y_test_batch = y_test_batch.to(device)
            
            # Project, spherically normalize, and append bias
            h_test_batch = torch.relu(torch.matmul(x_test_batch, projection_matrix))
            h_test_batch = torch.nn.functional.normalize(h_test_batch, p=2, dim=1)
            bias_test = torch.ones(h_test_batch.size(0), 1, device=device)
            X_test_batch = torch.cat((h_test_batch, bias_test), dim=1)
            
            # Inference via completed parameter matrix
            predictions = torch.matmul(X_test_batch, W)
            predicted_classes = torch.argmax(predictions, dim=1)
            
            correct += (predicted_classes == y_test_batch).sum().item()
            total += y_test_batch.size(0)

    accuracy = (correct / total) * 100
    print(f"--------------------------")
    print(f"Test Accuracy: {accuracy:.2f}%")
    print(f"Memory optimization active: Large matrices collected stream-wise.")

if __name__ == '__main__':
    main()