import torch
import mlflow
from urllib.parse import urlparse
import torch.nn as nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import datetime
from torch.autograd import Variable


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


mlflow.set_tracking_uri("http://172.17.0.1:5000")
dataset = pd.read_csv('datasets/train_set.csv')
testset = pd.read_csv('datasets/test_set.csv')

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)


criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)

for epoch in range(EPOCHS):
    inputs = Variable(torch.from_numpy(x_train))
    labels = Variable(torch.from_numpy(y_train))

    optimizer.zero_grad()

    outputs = model(inputs)

    loss = criterion(outputs, labels)
    print(loss)
    loss.backward()

    optimizer.step()

    print('epoch {}, loss {}'.format(epoch, loss.item()))

torch.save(model.state_dict(), 'model.pt')

with mlflow.start_run():
    test_input = x_train[0]

    mlflow.log_param("train size", dataset.size)
    mlflow.log_param("test size", testset.size)
    mlflow.log_param("epochs", EPOCHS)

    predicted = model(Variable(torch.from_numpy(x_train))).data.numpy()

    signature = mlflow.models.signature.infer_signature(
        x_train, predicted)

    mlflow.set_experiment("s434700")
    tracking_url_type_store = urlparse(mlflow.get_tracking_uri()).scheme
    if tracking_url_type_store != "file":
        mlflow.pytorch.log_model(model, "cs_model", registered_model_name="s434700", signature=signature,
                                 input_example=test_input)
    else:
        mlflow.pytorch.log_model(model, "cs_model",
                                 signature=signature, input_example=test_input)
        mlflow.pytorch.save_model(
            model, "cs_model", signature=signature, input_example=test_input)