import numpy as np
import pandas as pd
import torch
import torch.nn as tnn

filedir = '/home/ubuntu/Pulpit/TAU/mieszkania5'
#train size
learningRate = 0.0001 
epochs = 10000
#treainfile
trainfile = filedir + '/train/train.tsv'
#data files
dev0in = filedir + '/dev-0/in.tsv'
dev0out = filedir + '/dev-0/out.tsv' 
testAin = filedir + '/test-A/in.tsv'
testAout = filedir + '/test-A/out.tsv'


class linearRegression(tnn.Module):
    def __init__(self, dim_i, dim_o):
        super(linearRegression, self).__init__()
        self.linear = tnn.Linear(dim_i, dim_o)

    def forward(self, x):
        out = self.linear(x)
        return out

model = linearRegression(1, 1)
device = torch.device('cpu')
model.to(device)
print('model regresji gotowy')

#dane do treningu
trainfile_read = pd.raed_csv(trainfile, sep='\t', header=None, index_col=None)
train_data_sizes = np.array(trainfile_read[8].tolist(), dtype=np.float32).reshape(-1, 1)
train_data_prices = np.array(trainfile_read[0].tolist(), dtype=np.float32).reshape(-1, 1)

#dane do przewidywania
devfile_read = pd.raed_csv(dev0in, sep='\t', header=None, index_col=None)
testfile_in = pd.raed_csv(testAin, sep='\t', header=None, index_col=None)
dev_data_sizes = np.array(devfile_read[7].tolist(), dtype=np.float32)
test_data_sizes = np.array(testfile_in[7].tolist(), dtype=np.float32)
criterion = tnn.MSELoss() 
optimizer = torch.optim.SGD(model.parameters(), lr=learningRate)

print('dane wczytane')

#trening
for epoch in range(epochs):
    inputs = torch.from_numpy(train_data_sizes)
    labels = torch.from_numpy(train_data_prices)
    
    optimizer.zero_grad()

    outputs = model(inputs)

    loss = criterion(outputs, labels)

    loss.backward()

    optimizer.step()
print('trening zakonczony')
#przewidywanie ceny

predicted = model(torch.from_numpy(dev_data_sizes).requires_grad_()).data.numpy()

print('zapisywanie wyników dev-0')

dev_of = open(dev0out, 'w')
for i in predicted:
    dev_of.write(str(i[0])+'\n')
dev_of.close()

predicted = model(torch.from_numpy(test_data_sizes).requires_grad_()).data.numpy()

print('zapisywanie wyników test-A')

test_of = open(testAout, 'w')
for i in predicted:
    test_of.write(str(i[0])+'\n')
test_of.close()