#!/usr/bin/env python3

import sys
import torch

# Preprocessing i wektoryzację tekstów wydzialamy do osobnego modułu,
# z którego będzie korzystał zarówno kod do uczenia, jak i predykcji.
from analyzer import vectorizer, vector_length, process_line, vectorize_text

# Nasz model to zestaw wag w. Na koniec dojdziemy do modeli
# Transformer o wielu milionach wag; w pewnym sensie algorytm
# uczenia się nie zmieni.
# W naszym zadaniu lepiej zacząć od zerowych wag zamiast losowych.
# Znaczenie `requires_grad` zostanie omówione w kolejnym skrypcie.
w = torch.zeros(vector_length, dtype=torch.double, requires_grad=False)

# Hiperparametr uczenia
learning_rate = torch.tensor(.0032, dtype=torch.double)


def model(w, x):
    # @ to iloczyn wektorów/macierzy w PyTorchu
    return x @ w


# Funkcja kosztu.
def loss_fun(y_hat, y_exp):
    return (y_hat - y_exp)**2


# Co ile kroków będziemy wypisywali informacje o średniej funkcji kosztu.
# To nie jest hiperparametr uczenia, nie ma to żadnego, ani pozytywnego, ani
# negatywnego wpływu na uczenie.
step = 5000
i = 1
closs = torch.tensor(0.0, dtype=torch.double, requires_grad=False)

for line in sys.stdin:
    content, y_exp = process_line(line)

    x = vectorize_text(content)

    # wartość z predykcji
    y_hat = model(w, x)

    # wyliczamy funkcję kosztu
    loss = loss_fun(y_hat, y_exp)

    # Obliczamy gradient i aktualizujemy wagi.
    # TU SIĘ ODBYWA UCZENIE!
    grad = (y_hat - y_exp)
    w = w - learning_rate * x * grad

    # (Niedogodność tej wersji: jeśli zmienimy model lub funkcję kosztu,
    # będziemy musieli zróżniczkować wszystko i ustalić formułę na grad.

    # za jakiś czas pokazujemy uśrednioną funkcję kosztu
    closs += loss
    if i % step == 0:
        print("Sample item: ", y_exp.item(), " => ", y_hat.item(),
              " | Avg loss: ", (closs / step).item())
        closs = torch.tensor(0.0, dtype=torch.double, requires_grad=False)

    i += 1


# serializujemy nasz model
torch.save(w, "model.bin")