import torch


class TwoLayerNet(torch.nn.Module):
    def __init__(self, D_in, H, D_out):
        """
        In the constructor we instantiate two nn.Linear modules and assign them as
        member variables.
        """
        super(TwoLayerNet, self).__init__()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    def forward(self, x):
        """
        In the forward function we accept a Tensor of input data and we must return
        a Tensor of output data. We can use Modules defined in the constructor as
        well as arbitrary operators on Tensors.
        """
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

dtype = torch.float
device = torch.device("cuda:0")
# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 2140, 9216, 100, 30
def train(model,x,y,criterion,optimizer):
                """
       train NN
        """
    model.train()
    y_pred = model(x)
    loss = criterion(y_pred, y)
    print('train-loss',t, loss.item(),end=' ')
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()
def valid(model,x_valid,y_valid,criterion):
  """
       test NN
        """
    model.eval()
    y_pred = model(x_valid)
    loss = criterion(y_pred, y_valid)
    print('test-loss',t, loss.item(),end=' ')
    return loss.item()
# Create random Tensors to hold inputs and outputs
x_train =  torch.tensor(torch.from_numpy(X_train),device=device,dtype=dtype)
y_train =   torch.tensor(torch.from_numpy(y_train),device=device,dtype=dtype)
x_valid =  torch.tensor(torch.from_numpy(X_valid),device=device,dtype=dtype)
y_valid =   torch.tensor(torch.from_numpy(y_valid),device=device,dtype=dtype)
model = TwoLayerNet(D_in, H, D_out).to(device)
loss_train=[]
loss_valid=[]
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
for t in range(800):
    loss_train.append(train(model,x_train,y_train,criterion,optimizer))
    loss_valid.append(valid(model,x_valid,y_valid,criterion))
    print()