#  pytorch_mlp_2class.py
#  G. Cowan / RHUL Physics / June 2025
#  Simple program to illustrate classification with pytorch
#  using two-class network output.
#  Based on simpleClassifier_mlp.py which used scikit-learn

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from sklearn.model_selection import train_test_split
from sklearn import metrics


#  read the data in from files, 
#  assign target values 1 for signal, 0 for background
sigData = np.loadtxt('signal.txt')
nSig = sigData.shape[0]
sigTargets = np.ones(nSig, dtype=np.int64)

bkgData = np.loadtxt('background.txt')
nBkg = bkgData.shape[0]
bkgTargets = np.zeros(nBkg, dtype=np.int64)

# Concatenate arrays into data X and targets y
X_np = np.concatenate((sigData,bkgData),0)
X_np = X_np[:,0:2]                    # at first, only use x1 and x2
y_np = np.concatenate((sigTargets, bkgTargets))
X = torch.from_numpy(X_np).float()     # float32 features
y = torch.from_numpy(y_np).long()      # int64 class labels

# Use scikit-learn to split into train/test sets
X_train_np, X_test_np, y_train_np, y_test_np = train_test_split( X_np, y_np,
    test_size=0.2, random_state=42, stratify=y_np)
X_train = torch.from_numpy(X_train_np).to(X.dtype)
y_train = torch.from_numpy(y_train_np).to(y.dtype)
X_test = torch.from_numpy(X_test_np).to(X.dtype)
y_test = torch.from_numpy(y_test_np).to(y.dtype)

# Define model
model = nn.Sequential(
    nn.Linear(2, 10),
    nn.ReLU(),
    nn.Linear(10, 2)
)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Training loop
for epoch in range(200):
    optimizer.zero_grad()
    output = model(X_train)
    loss = criterion(output, y_train)
    loss.backward()
    optimizer.step()
    
    if epoch % 10 == 0:
        output_train = model(X_train)
        pred_train = output_train.argmax(dim=1)
        acc_train = (pred_train == y_train).float().mean()
        output_test = model(X_test)
        pred_test = output_test.argmax(dim=1)
        acc_test = (pred_test == y_test).float().mean()
        print(f"Epoch {epoch:3d} | Loss: {loss.item():.4f} | ",
              f"Accuracy (train): {acc_train:.4f} | ",
              f"Accuracy (test) = {acc_test:.4f}")


# make a scatter plot
fig, ax = plt.subplots(1,1)
plt.gcf().subplots_adjust(bottom=0.15)
plt.gcf().subplots_adjust(left=0.15)
ax.set_xlim((-2.5,3.5))
ax.set_ylim((-2,4))
x0,x1 = ax.get_xlim()
y0,y1 = ax.get_ylim()
ax.set_aspect(abs(x1-x0)/abs(y1-y0))       # make square plot
xtick_spacing = 0.5
ytick_spacing = 2.0
ax.yaxis.set_major_locator(ticker.MultipleLocator(xtick_spacing))
ax.yaxis.set_major_locator(ticker.MultipleLocator(ytick_spacing))
plt.scatter(sigData[:,0], sigData[:,1], s=3, color='dodgerblue', marker='o')
plt.scatter(bkgData[:,0], bkgData[:,1], s=3, color='red', marker='o')

# add decision boundary to scatter plot
x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5
y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5
h = 0.01  # step size in the mesh
xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))

grid = torch.tensor(np.c_[xx.ravel(), yy.ravel()], dtype=torch.float32)
with torch.no_grad():
    logits = model(grid)
    t_values = logits[:, 1] - logits[:, 0]  # Decision function
t_values_np = t_values.reshape(xx.shape).numpy()

plt.contour(xx, yy, t_values_np, 1, colors='k')
plt.xlabel(r'$x_{1}$', labelpad=0)
plt.ylabel(r'$x_{2}$', labelpad=15)
plt.savefig("scatterplot.pdf", format='pdf')
plt.show()