import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from ionic_ml_simulator import IonicMLSimulator
import numpy as np

# Load MNIST subset
X, y = fetch_openml('mnist_784', version=1, return_X_y=True, as_frame=False)
X = X[:10000] / 255.0
y = y[:10000].astype(int)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# Define a tiny model
model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
loss_fn = nn.CrossEntropyLoss()
sim = IonicMLSimulator(d_model=128, learning_rate=0.01)

# Convert data to list of (x,y) tuples
train_data = list(zip(X_train, y_train))

# Training loop (classical, not O(1))
"""for epoch in range(10):
    for i in range(0, len(train_data), 32):
        batch = train_data[i:i+32]
        loss = sim.train_step(batch, model, loss_fn)
    print(f"Epoch {epoch}, loss {loss:.4f}")"""

for epoch in range(1000):
    idx = np.random.randint(0,8000,100)
    X0 = X_train[idx]
    y0 = y_train[idx]
    batch0 = list(zip(X0, y0))
    for _ in range(100):
        idx = np.random.randint(0,8000,100)
        X = X_train[idx]
        yt = y_train[idx]
        batch = list(zip(X, yt))
        loss = sim.train_step(batch0, model, loss_fn)
        loss = sim.train_step(batch, model, loss_fn)
    print(f"Epoch {epoch}, loss {loss:.4f}")
    if epoch % 1==0:
        # Model evaluation / Testing
        model.eval()
        test_loss = 0
        correct = 0
        with torch.no_grad():
            for x, y in zip(X_test, y_test):
                pred = model(torch.tensor(x, dtype=torch.float32))
                loss = loss_fn(pred.unsqueeze(0), torch.tensor([y]))
                test_loss += loss.item()
                correct += (pred.argmax(0) == y).type(torch.float).item()

        test_loss /= len(X_test)
        accuracy = correct / len(X_test)
        print(f"\nTest Evaluation:")
        print(f"Test Accuracy: {accuracy * 100:.2f}%")
        print(f"Test Avg Loss: {test_loss:.4f}")