"""
    "Building a Simple MLP from Scratch Using PyTorch" by Aymen Noor
    https://medium.com/@mn05052002/building-a-simple-mlp-from-scratch-using-pytorch-7d50ca66512b
"""
import torch
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

class SimpleMLP(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleMLP, self).__init__()
        self.W1 = torch.randn(input_size, hidden_size, requires_grad=True)
        self.b1 = torch.randn(1, hidden_size, requires_grad=True)
        self.W2 = torch.randn(hidden_size, output_size, requires_grad=True)
        self.b2 = torch.randn(1, output_size, requires_grad=True)

    def forward(self, X):
        self.z1 = torch.matmul(X, self.W1) + self.b1
        self.a1 = torch.sigmoid(self.z1)  # Hidden layer activation
        self.z2 = torch.matmul(self.a1, self.W2) + self.b2
        self.a2 = torch.sigmoid(self.z2)  # Output layer activation
        return self.a2

    def backward(self,X,y,output,lr=0.01):
        m=X.shape[0]
        dz2=output-y
        dW2=torch.matmul(self.a1.T,dz2)
        db2=torch.sum(dz2,axis=0)/m

        da1=torch.matmul(dz2,self.W2.T)
        dz1=da1*(self.a1*(1-self.a1))
        dw1=torch.matmul(X.T,dz1)/m
        db1 = torch.sum(dz1, axis=0) / m
      
        with torch.no_grad():
            self.W1 -= lr * dw1
            self.b1 -= lr * db1
            self.W2 -= lr * dW2
            self.b2 -= lr * db2

    def train(self, X, y, epochs=1000, lr=0.01):
        losses = []
        for epoch in range(epochs):
            output = self.forward(X)
            #Compute loss using (Mean Squared Error)
            loss = torch.mean((output - y) ** 2)
            losses.append(loss.item())
            #update weights
            self.backward(X, y, output, lr)
            if (epoch + 1) % 100 == 0:
                print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")
        return losses
    
    
if __name__ == '__main__':    
    
    # Generate dataset
    X, y = make_moons(n_samples=1000, noise=0.2, random_state=42)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # Standardize the data
    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)
    
    # Convert to PyTorch tensors
    X_train = torch.tensor(X_train, dtype=torch.float32)
    y_train = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1)
    X_test = torch.tensor(X_test, dtype=torch.float32)
    y_test = torch.tensor(y_test, dtype=torch.float32).reshape(-1, 1)
    
    input_size = 2
    hidden_size = 4
    output_size = 1
    model = SimpleMLP(input_size, hidden_size, output_size)
    
    #Train  model and store the losses
    losses = model.train(X_train, y_train, epochs=1000, lr=0.1)
    
    plt.plot(losses)
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Training Loss over Epochs")
    plt.show()
    
    with torch.no_grad():
        test_output = model.forward(X_test)
        test_output = (test_output > 0.5).float() 
    accuracy = torch.mean((test_output == y_test).float())
    print(f"Test Accuracy: {accuracy.item() * 100:.2f}%")
    
    print('finished')