#!/usr/bin/env python3
import argparse

import numpy as np
import sklearn.datasets
import sklearn.metrics
import sklearn.model_selection

parser = argparse.ArgumentParser()

parser.add_argument("--learning_rate", default=0.01, type=float, help="Learning rate")
parser.add_argument("--epochs", default=50, type=int, help="Počet epoch na trénování Minibatch SGD")
parser.add_argument("--batch_size", default=10, type=int, help="Velikost batche")
parser.add_argument("--data_size", default=100, type=int, help="Velikost datasetu")
parser.add_argument("--test_size", default=0.5, type=float, help="Velikost testovací množiny")
parser.add_argument("--seed", default=42, type=int, help="Náhodný seed")

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def make_prediction(dato, weights):
    prediction = 0
    for d, w in zip(dato, weights):
        prediction += d*w
    return sigmoid(prediction)

def make_predictions(data, weights):
    arr = []
    for dato in data:
        arr.append(make_prediction(dato, weights))
    return arr

def make_classifications(data, weights):
    arr = []
    for dato in data:
        arr.append(1 if make_prediction(dato, weights) > 0.5 else 0)
    return arr

def main(args: argparse.Namespace):
    # Nastavení seedu generátoru náhodných čísel
    generator = np.random.RandomState(args.seed)

    # Vytvoření náhodného datasetu na klasifikační úlohu se dvěma třídami
    # třídy mají hodnotu 0 a 1
    data, target = sklearn.datasets.make_classification(n_samples=args.data_size, random_state=args.seed)

    # Přidání sloupce jedniček pro bias
    data = np.concatenate([data, np.ones([args.data_size, 1])], axis=1)

    # TODO: (SGD) Rozdělte dataset na trénovací a testovací část, funkci z knihovny sklearn
    # předejte argumenty `test_size=args.test_size, random_state=args.seed`.
    train_data, test_data, train_target, test_target = \
        sklearn.model_selection.train_test_split(data, target, test_size=args.test_size, random_state=args.seed)

    # Vytvoření náhodných vah
    weights = generator.uniform(size=train_data.shape[1], low=-0.1, high=0.1)

    # Zjištění počtu prvků v trénovací množině.
    # `train_data` je 2D pole, kde první dimenze je počet prvků a druhá dimenze je počet featur.
    # Ekvivalentně šlo zjistit počet prvků i pomocí:
    # `train_size = int(args.data_size * (1 - args.test_size))`
    train_size = train_data.shape[0]

    for epoch in range(args.epochs):
        # TODO: (SGD) Pro každou epochu náhodně zamíchejte trénovací množinu.
        # Na zamíchání použijte generator.permutation, které vrátí náhodnou permutaci
        # čísel od 0 do počtu prvků v množině.
        # Nemíchejte původní trénovací dataset, ale jen jeho kopii, tedy
        # `train_data` zůstane nezměněný přes všechny epochy.
        permutation = generator.permutation(train_size)

        # TODO: Pro každou batch s velikostí args.batch_size spočítejte gradient
        # a upravte váhy.
        # Můžete předpokládat, že args.batch_size je dělitelem počtu prvků v trénovací množině.
        for batch_start_idx in range(0, train_size, args.batch_size):
            gradient = np.zeros(train_data.shape[1])
            # Procházím batch který začíná na indexu `batch_start_idx` a končí na indexu `batch_start_idx+args.batch_size-1`.
            # Slicing v Pythonu `permuration[...:...]` má polootevřený interval, tedy poslední index není zahrnut.
            for batch_elem_idx in permutation[batch_start_idx:batch_start_idx+args.batch_size]:
                dato = train_data[batch_elem_idx]
                dato_target = train_target[batch_elem_idx]
                prediction = make_prediction(dato, weights)
                for i in range(train_data.shape[1]):
                    gradient[i] += (prediction - dato_target) * dato[i]

            weights = weights - args.learning_rate * gradient / args.batch_size

            # Bez numpy:
            # for i in range(train_data.shape[1]):
            #     weights[i] = weights[i] - args.learning_rate * gradient[i] / args.batch_size

        # TODO: Na konci každé epochy spočítejte accuracy metriku na trénovací a testovací množině.
        # Accuracy metrika se počítá jako počet správných predikcí děleno počtem všech predikcí.
        # Metriku můžete spočítat explicitně nebo pomocí funkce `sklearn.metrics.accuracy_score`.
        # Metrika accuracy chce na vstupu již klasifikovaná data, ne pravděpodobnosti.
        train_accuracy = sklearn.metrics.accuracy_score(train_target, make_classifications(train_data, weights))
        test_accuracy = sklearn.metrics.accuracy_score(test_target, make_classifications(test_data, weights))

        # TODO: Poté spočítejte chybu logistické regrese na trénovací a testovací množině.
        # Chybu můžete spočítat explicitně nebo využít funkce `log_loss` v knihovně sklearn.
        # log-loss namísto accuracy chce na vstupu pravděpodobnosti, ne již klasifikovaná data.
        train_log_loss = sklearn.metrics.log_loss(train_target, make_predictions(train_data, weights))
        test_log_loss = sklearn.metrics.log_loss(test_target, make_predictions(test_data, weights))
        print(f"Epoch {epoch+1}: train loss {train_log_loss:.6f} acc {train_accuracy*100:.2f}, test loss {test_log_loss:.6f} acc {test_accuracy*100:.2f}")

if __name__ == "__main__":
    args = parser.parse_args()
    main(args)