#include <bits/stdc++.h>

using namespace std;

// Ograniczamy sie do cyfr [0, 1, ..., NUM_CLASSES-1].
const int NUM_CLASSES = 5;

// Obrazek (jako ciag 28 * 28 liczb) + klasa.
struct Datapoint {
    vector <int> image;
    int label;
};

unsigned char readChar(ifstream &file) {
    // Wczytujemy z pliku 1 bajt.
    char c;
    file.read(&c, 1);

    return c;
}

int readInt(ifstream &file) {
    // Wczytujemy z pliku 4 bajty, i skladamy je w int'a.
    int result = 0;

    for (int byte = 0; byte < 4; byte++) {
        result = (result <<= 8) + readChar(file);
    }

    return result;
}

vector <vector<int>> read_mnist_images(string filename) {
    ifstream file(filename, ios::binary);

    // Pierwsza liczba w pliku to suma kontrolna.
    const int CHECKSUM_EXPECTED = 2051;

    int checksum = readInt(file);
    assert(checksum == CHECKSUM_EXPECTED);

    int number_of_images = readInt(file);
    int number_of_rows = readInt(file);
    int number_of_columns = readInt(file);
    int image_size = number_of_rows * number_of_columns;

    vector <vector<int>> images(number_of_images, vector <int> (image_size));
    for (auto &image : images) {
        for (int &pixel : image) {
            pixel = readChar(file);
        }
    }

    return images;
}

vector <int> read_mnist_labels(string filename) {
    ifstream file(filename, ios::binary);

    // Pierwsza liczba w pliku to suma kontrolna.
    const int CHECKSUM_EXPECTED = 2049;

    int checksum = readInt(file);
    assert(checksum == CHECKSUM_EXPECTED);

    int number_of_labels = readInt(file);

    vector <int> labels(number_of_labels);
    for (int &label : labels) {
        label = readChar(file);
    }

    return labels;
}

vector <Datapoint> read_mnist_data(string filename_images, string filename_labels) {
    // Wczytujemy dane MNIST.
    auto images = read_mnist_images(filename_images);
    auto labels = read_mnist_labels(filename_labels);

    vector <Datapoint> data;
    for (int i = 0; i < (int) images.size(); i++) {
        if (labels[i] >= NUM_CLASSES) {
            continue;
        }

        data.push_back({images[i], labels[i]});
    }

    return data;
}

void show_datapoint(Datapoint &datapoint) {
    // Prymitywne wyswietlanie obrazka na standardowe wyjscie.

    cout << "Label: " << datapoint.label << endl;
    for (int i = 0; i < (int) datapoint.image.size(); i++) {
        int pixel_value = datapoint.image[i];

        if (pixel_value > 100) {
            cout << '#';
        } else {
            cout << '.';
        }

        if (i % 28 == 27) {
            // Koniec wiersza.
            cout << '\n';
        }
    }
}

int predict(vector <int> &image, vector <vector <double>> &coefficients) {
    vector <double> sums(NUM_CLASSES, 0.0);

    // Liczymy odpowiednia sume dla kazdej mozliwej cyfry...
    for (int digit = 0; digit < NUM_CLASSES; digit++) {
        for (int i = 0; i < (int) image.size(); i++) {
            sums[digit] += image[i] * coefficients[digit][i];
        }
    }

    // ...wybieramy cyfre o najwiekszej sumie jako nasza predykcje.
    int max_digit = 0;
    for (int digit = 1; digit < NUM_CLASSES; digit++) {
        if (sums[digit] > sums[max_digit]) {
            max_digit = digit;
        }
    }

    return max_digit;
}

double test(vector <Datapoint> &test_data, vector <vector <double>> &coefficients) {
    int correct = 0;
    for (auto &datapoint : test_data) {
        // Przewidujemy klase obrazka testowego...
        int prediction = predict(datapoint.image, coefficients);

        // ...i sprawdzamy, czy sie udalo.
        correct += prediction == datapoint.label;
    }

    return (double) correct / test_data.size();
}

void test_and_print(vector <Datapoint> &test_data, vector <vector <double>> &coefficients) {
    auto correct_fraction = test(test_data, coefficients);
    cout << 100.0 * correct_fraction << "%" << endl;
}

double random_double() {
    // Zwraca losowa liczbe z zakresu [-1, 1].
    const int M = 100;

    // Losujemy z zakresu [0, 1]...
    double value = (double) (rand() % (M + 1)) / M;

    // ...i przeskalowujemy do [-1, 1].
    return 2 * value - 1;
}

vector <vector <double>> train(vector <Datapoint> &train_data, int num_steps) {
    int image_length = train_data[0].image.size(); // = 784

    // Inicjalizujemy wspolczynniki na 0.
    vector <vector <double>> coefficients(NUM_CLASSES, vector <double> (image_length, 0.0));

    cout << "Wynik na danych treningowych bez treningu: ";
    test_and_print(train_data, coefficients);

    double score = test(train_data, coefficients);

    // Wykonujemy num_steps krokow...
    for (int step = 0; step < num_steps; step++) {
        auto new_coefficients = coefficients;

        // ...w kazdym kroku zmieniamy losowo wspolczynniki...
        for (int digit = 0; digit < NUM_CLASSES; digit++) {
            for (int i = 0; i < image_length; i++) {
                new_coefficients[digit][i] += random_double();
            }
        }

        // ...liczymy wynik po zmianie, jeśli jest lepszy, to zapisujemy.
        double new_score = test(train_data, new_coefficients);
        if (new_score > score) {
            coefficients = new_coefficients; score = new_score;
            cout << "Krok " << step << ": " << new_score << endl;
        }
    }

    cout << "Wynik na danych treningowych po treningu: ";
    test_and_print(train_data, coefficients);

    return coefficients;
}

int main() {
    srand(0);

    // Ladujemy dane treningowe.
    auto train_data = read_mnist_data("train-images-idx3-ubyte", "train-labels-idx1-ubyte");

    // Wyswietlamy dwa pierwsze obrazki.
    show_datapoint(train_data[0]);
    show_datapoint(train_data[1]);

    cout << "Wczytano " << train_data.size() << " obrazkow treningowych." << endl;

    // Do treningu bierzemy tylko 10000 obrazków. Moglibysmy wziac wszystkie, ale
    // 10k wystarcza by dostac dobre wyniki, a trening bedzie dzialac szybciej.
    random_shuffle(train_data.begin(), train_data.end());
    train_data.resize(10000);

    auto coefficients = train(train_data, 5000);

    // Ladujemy dane testowe.
    auto test_data = read_mnist_data("t10k-images-idx3-ubyte", "t10k-labels-idx1-ubyte");

    cout << "Wynik na danych testowych: ";
    test_and_print(test_data, coefficients);

    return 0;
}