from sklearn.model_selection import train_test_split
import numpy as np
import torch
from torch import nn
from torch.optim import lr_scheduler
from sklearn.metrics import mean_squared_error
import sys

n = len( sys.argv )
if n != 3 :
	print( "Args: seed lr" )
	sys.exit(1)

torch.manual_seed( int(sys.argv[1]) )
learning_rate = float( sys.argv[2] )


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

X = A[:,0:2]
yi = A[:,2]
y = np.where(yi==0, -1, 1 )

print( y )

# 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=1 )

dtype = torch.float32

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

X_prueba = torch.tensor( X_prueba, dtype=dtype )
y_prueba = y_prueba.reshape(-1,1)
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.Tanh(),
			nn.Linear(2, 1),
			# La salida es la función identidad para un regresor
		)

	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.MSELoss( )
# loss_fn = nn.MSELoss( reduction='sum')
# loss_fn = nn.CrossEntropyLoss( )

batch_size = 10

optimizador = torch.optim.SGD(modelo.parameters(), lr=learning_rate,
                              momentum=0.9, nesterov=True )

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

# scheduler = lr_scheduler.StepLR(optimizador, step_size=4, gamma=0.5)


def train_loop( X, y, model, loss_fn, optimizer):
	size, _ = X.shape 
	print( size )
	# 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
		optimizer.zero_grad()
		loss.backward()
		optimizer.step( )

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

		i += batch_size

	# scheduler.step()

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)

	mse = mean_squared_error( y, pred )

	print( "Error MSE de entrenamiento:", mse ) 
	return mse
	

epochs = 1000
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( "LR:", learning_rate )
print("Done!")

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

