from sklearn.model_selection import train_test_split
import numpy as np
import torch
from torch import nn
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score

torch.manual_seed(2)

# Leo los datos
A = np.loadtxt( "d.txt" )
print( A.shape )

X = A[:,0:2]
y = A[:,2]

# Dividimos los datos en datos de entrenamiento
# y datos de prueba
X_entrena, X_prueba, y_entrena, y_prueba = train_test_split( X, y, random_state=3 )

y_entrena = y_entrena.reshape( -1,1 )
y_prueba = y_prueba.reshape( -1,1 )

dtype = torch.float32

# Cambiamos los datos a tensor
X_entrena = torch.tensor( X_entrena, dtype=dtype )
y_entrena = torch.tensor( y_entrena, dtype=dtype )

X_prueba = torch.tensor( X_prueba, dtype=dtype )
y_prueba = torch.tensor( y_prueba, dtype=dtype )


# Tenemos que crear el modelo de la red neuronal

class RedNeuronal_2( nn.Module ):
	def __init__(self):
		super().__init__()
		self.arquitectura = nn.Sequential(
			nn.Linear(2, 2),
			nn.Sigmoid(),
			nn.Linear(2, 1),
			nn.Sigmoid()
		)

	def forward(self, x):
		r = self.arquitectura(x)
		return r

# Instanciamos el modelo
modelo = RedNeuronal_2()
print(modelo)


# Initialize the loss function
loss_fn = nn.BCELoss()

learning_rate = 0.19
batch_size = 10

optimizador = torch.optim.Adam(modelo.parameters(), lr=learning_rate)


def train_loop( X, y, model, loss_fn, optimizer):
	size, _ = X.shape 
	# Set the model to training mode - important for batch normalization and dropout layers
	# Unnecessary in this situation but added for best practices
	model.train()
	i = 0
	while i < size :
		Xe = X[i:i+batch_size,:]
		ye = y[i:i+batch_size,:]

		# Compute prediction and loss
		pred = model(Xe)
		loss = loss_fn(pred, ye)

		# Backpropagation
		loss.backward()
		optimizer.step()
		optimizer.zero_grad()

		# if i % 10 == 0:
		loss, current = loss.item(), i+batch_size
		print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

		i += batch_size

def test_loop( X, y, model, loss_fn):
	# Set the model to evaluation mode - important for batch normalization and dropout layers
	# Unnecessary in this situation but added for best practices
	model.eval()

	# Evaluating the model with torch.no_grad() ensures that no gradients are computed during test mode
	# also serves to reduce unnecessary gradient computations and memory usage for tensors with requires_grad=True

	with torch.no_grad() :
		pred = model(X)
		salida = torch.where( pred >= 0.5, 1, 0 )

	exactitud = accuracy_score( y, salida )

	print( "Exactitud de ntrenamiento:", exactitud ) 
	

epochs = 50
for t in range(epochs):
	print(f"Epoch {t+1}\n-------------------------------")
	train_loop( X_entrena, y_entrena, modelo, loss_fn, optimizador)
	test_loop( X_prueba, y_prueba, modelo, loss_fn)
print("Done!")

torch.save(modelo, 'modelo0.pth')

