# Always suppress warnings for cleaner output
import warnings; warnings.filterwarnings("ignore")
import numpy as np
np.random.seed(42)

# Data setup
from sklearn.datasets import load_digits
digits = load_digits()
X, y = digits.data, digits.target
print(X.shape)
print(y.shape)

# Let us print and view the first image
import matplotlib.pyplot as plt
plt.gray()
plt.matshow(digits.images[0])
plt.title(f"True Label: {digits.target[0]}")
plt.show()

# Let us check the distribution of digit labels
import pandas as pd
import seaborn as sns
sns.set_style("whitegrid")
df_labels = pd.DataFrame({'label': y})
sns.countplot(x='label', data=df_labels, palette='viridis')
plt.title('Number of Samples per Digit')
plt.show()

# Prepare data for PyTorch: train test split and normalization
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
print(X_train_scaled.shape, X_test_scaled.shape)

# Convert data to torch tensors and wrap in datasets
import torch
from torch.utils.data import TensorDataset, DataLoader
X_train_tensor = torch.tensor(X_train_scaled, dtype=torch.float32)
X_test_tensor = torch.tensor(X_test_scaled, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.long)
y_test_tensor = torch.tensor(y_test, dtype=torch.long)
train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32)

# Define the logistic regression model in PyTorch
import torch.nn as nn
class LogisticRegressionModel(nn.Module):
    def __init__(self):
        super(LogisticRegressionModel, self).__init__()
        self.linear = nn.Linear(64, 10)
    def forward(self, x):
        outputs = self.linear(x)
        return outputs
model = LogisticRegressionModel()

# Set up loss function and optimizer
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Run the training loop
epochs = 20
for epoch in range(epochs):
    model.train()
    total_loss = 0
    for batch_x, batch_y in train_loader:
        optimizer.zero_grad()
        outputs = model(batch_x)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * batch_x.size(0)
    avg_loss = total_loss / len(train_loader.dataset)
    if (epoch + 1) % 5 == 0 or epoch == 0:
        print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")

# Evaluate test accuracy after training
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for batch_x, batch_y in test_loader:
        outputs = model(batch_x)
        _, predicted = torch.max(outputs, 1)
        correct += (predicted == batch_y).sum().item()
        total += batch_y.size(0)
accuracy = correct / total * 100
print(f"Test Accuracy: {accuracy:.2f}%")

# Show predictions for some test images
import numpy as np
examples = np.random.choice(len(X_test_tensor), 8, replace=False)
model.eval()
with torch.no_grad():
    for idx in examples:
        img = X_test_tensor[idx].numpy().reshape(8, 8)
        out = model(X_test_tensor[idx].unsqueeze(0))
        pred = out.argmax(dim=1).item()
        plt.imshow(img, cmap='gray')
        plt.title(f"Predicted: {pred}, True: {y_test[idx]}")
        plt.axis('off')
        plt.show()

# Look at a confusion matrix for test results
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
all_preds = []
all_true = []
model.eval()
with torch.no_grad():
    for batch_x, batch_y in test_loader:
        outputs = model(batch_x)
        _, predicted = torch.max(outputs, 1)
        all_preds.extend(predicted.numpy())
        all_true.extend(batch_y.numpy())
cm = confusion_matrix(all_true, all_preds)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=np.arange(10))
disp.plot(cmap='Blues')
plt.title('Confusion Matrix: Test Set')
plt.show()

# Optional: Type a digit image as numbers for prediction
import numpy as np
user_pixels = []
print("Enter 64 pixel values (0 16) for an 8x8 image, separated by space:")
for i in range(8):
    row_values = input(f"Row {i+1}/8: ").split()
    user_pixels.extend([int(x) for x in row_values])
user_img = np.array(user_pixels).reshape(1, -1)
user_img_scaled = scaler.transform(user_img)
user_tensor = torch.tensor(user_img_scaled, dtype=torch.float32)
with torch.no_grad():
    output = model(user_tensor)
    predicted = output.argmax(dim=1).item()
print(f"Predicted digit: {predicted}")