"""
    Create and Train Simple MLP for scikit make_moons.
    Convert and Save in OpenVino Intermediate Representation (IR)

"""

import openvino as ov
import torch
from torch import nn
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(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.layer1 = nn.Linear(2, 4)   # Hidden layer
        self.layer2 = nn.Linear(4, 1)   # Output layer

    def forward(self, x):
        x = self.layer1(x)
        x = torch.sigmoid(x)  # Activation function
        x = self.layer2(x)
        return x

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)

    # Plot make_moons
    fig,ax1 = plt.subplots(nrows=1,ncols=1,figsize=(8,4))
    ax1.scatter(X[:,0],X[:,1],c=y)
    plt.show() 
    
    # 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)
    
    model = SimpleMLP()
    criterion = nn.MSELoss()  # Mean Squared Error for regression
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # Stochastic Gradient Descent

    losses = []
    for epoch in range(1000):  # Number of epochs
        optimizer.zero_grad()  # Clear gradients
        outputs = model(X_train)  # Forward pass
        loss = criterion(outputs, y_train)  # Compute loss
        loss.backward()  # Backward pass
        losses.append(loss.item())
        optimizer.step()  # Update parameters
  
    plt.plot(losses)
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title("Training Loss over Epochs")
    plt.show()
    
    with torch.no_grad():
        raw_output = model.forward(X_test)
        test_output = (raw_output > 0.5).float() 
    accuracy = torch.mean((test_output == y_test).float())
    print(f"Test Accuracy: {accuracy.item() * 100:.2f}%")

    # Convert and Save to OpenVino Intermediate Representation (IR)
    ov_input=(ov.PartialShape([1, 2]), ov.Type.f32)
    ov_model = ov.convert_model(model, input=ov_input)
    compiled_model = ov.compile_model(ov_model, "AUTO")
    ov.save_model(ov_model,'./mlp.xml')

    print("finished")