#!/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("--classes", default=10, type=int, help="Počet tříd pro klasifikaci")
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 count_probabilities(data, weights):
    if data.ndim == 1:
        data = np.reshape(data, (1, -1))
    
    # Výpočet lineárních modelů (lineární kombinace vah a dat)
    predictions = data @ weights.T
    # Pro numerickou stabilitu výpočtu softmaxu
    predictions -= np.max(predictions)

    _sum = np.sum(np.exp(predictions), axis=1)

    return np.exp(predictions) / _sum[:, np.newaxis]

def make_predictions(data, weights):
    return np.argmax(count_probabilities(data, weights), axis=1)

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 s počtem tříd `args.classes`.
    # Třídy mají hodnotu 0 až args.classes-1
    data, target = sklearn.datasets.make_classification(
        n_samples=args.data_size, n_informative=args.classes,
        n_classes=args.classes, 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=[args.classes, 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([args.classes, 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.
            batch = permutation[batch_start_idx:batch_start_idx+args.batch_size]
            data = train_data[batch]
            targets = train_target[batch]
            probabilities = count_probabilities(data, weights)

            # Zakódování targetů do one-hot reprezentace
            one_hot_targets = np.arange(args.classes)[np.newaxis, :] == targets[:, np.newaxis]

            # Vypočítaní gradientu pro danou batch
            gradient = (probabilities - one_hot_targets).T @ data

            # Aktualizace vah
            weights = weights - args.learning_rate * gradient / 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_predictions(train_data, weights))
        test_accuracy = sklearn.metrics.accuracy_score(test_target, make_predictions(test_data, weights))
        
        # TODO: Poté spočítejte chybu softmaxu 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, count_probabilities(train_data, weights))
        test_log_loss = sklearn.metrics.log_loss(test_target, count_probabilities(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)