import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import seaborn as sns
import torch
import datetime
from torch.autograd import Variable
import csv

INPUT_DIM = 1
OUTPUT_DIM = 1
LEARNING_RATE = 0.01
EPOCHS = 100


dataset = pd.read_csv('datasets/train_set.csv')

# create dummy data for training
x_values = [datetime.datetime.strptime(
    item, "%Y-%m-%d").month for item in dataset['date'].values]
x_train = np.array(x_values, dtype=np.float32)
x_train = x_train.reshape(-1, 1)

y_values = [min(dataset['result_1'].values[i]/dataset['result_2'].values[i], dataset['result_2'].values[i] /
                dataset['result_1'].values[i]) for i in range(len(dataset['result_1'].values))]
y_train = np.array(y_values, dtype=np.float32)
y_train = y_train.reshape(-1, 1)


class LinearRegression(torch.nn.Module):
    def __init__(self, inputSize, outputSize):
        super(LinearRegression, self).__init__()
        self.linear = torch.nn.Linear(inputSize, outputSize)

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


model = LinearRegression(INPUT_DIM, OUTPUT_DIM)

model.load_state_dict(torch.load("model/model.pt"), strict=False)


# testing data

with torch.no_grad():  # we don't need gradients in the testing phase
    predicted = model(Variable(torch.from_numpy(x_train))).data.numpy()


with open('model_results.csv', mode='w') as filee:
    writer = csv.writer(filee, delimiter=',', quotechar='"',
                        quoting=csv.QUOTE_MINIMAL)

    writer.writerow(['x', 'y', 'predicted_y'])

    for i in range(len(x_train)):
        writer.writerow([x_train[i][0], y_train[i][0], predicted[i][0]])