// simple_backprop.cpp : Defines the entry point for the console application.
//
#include <cmath>
#include <iostream>

#include <fstream>

#include <iomanip>
#include <random>
#include <algorithm>
#include <cstdint>

#include <arpa/inet.h> // for ntohl

using namespace std;

typedef float flt;

//#pragma omp

/*
inline void activation_function(flt a, flt &result, flt &der) {
    result=a>-20? 1.0 / (1.0 + exp(-a)) : 0.0;
    der= result*(flt(1.0) - result);
}*/


inline void activation_function(flt a, flt &result, flt &der) {// der = dresult/da
    const float param=0.1;
    result = a>0.0 ? a : a*param;
    der = a>0.0 ? 1.0 : param;
}

flt loss_function(flt a, flt b) {
    return flt(0.5)*(a-b)*(a-b);
}

flt loss_function_der(flt a, flt b) {// d loss_function(a,b)/db
    return b - a;
}

const flt step_grow=1.2f;
const flt step_shrink=0.5;

const flt step_min=1E-10;
const flt step_max= 10.0;

template<class T>
T clamp(T a, T low, T high){
    return std::min(std::max(a,low),high);
}

std::random_device rdev;
mt19937 rng(rdev());
std::uniform_real_distribution<> r(0, 1);

template<int n_inputs_, int n_outputs_>
struct layer {
    static const int n_inputs=n_inputs_;
    static const int n_outputs=n_outputs_;
    flt input_values[n_inputs]{};

    flt weights[n_outputs][n_inputs]{};
    flt old_weights[n_outputs][n_inputs]{};

    flt derr_per_weight_sums[n_outputs][n_inputs]{};
    flt old_derr_per_weight_sums[n_outputs][n_inputs]{};

    flt weight_step_size[n_outputs][n_inputs]{};

    flt input_derivatives[n_inputs]{};
    flt output_values[n_outputs]{};
    flt dout_per_dsum[n_outputs]{};
    flt derr_per_dout[n_outputs]{};

    layer() {// initialize weights to sum to 1 for now


        for (int o = 0; o < n_outputs; ++o)
            for (int i = 0; i < n_inputs; ++i){
                weights[o][i] = old_weights[o][i] = r(rng)*0.1-0.05;//0.1*(r(rng)*0.2+0.9)/n_inputs;
                weight_step_size[o][i]=0.01;
            }
    }

    void forward(const flt inputs[n_inputs]) {
        //#pragma omp parallel_for
        for (int i = 0; i < n_inputs; ++i) {
            input_values[i] = inputs[i];
        }
        //#pragma omp parallel_for
        for (int o = 0; o < n_outputs; ++o) {
            flt sum = 0;
            for (int i = 0; i < n_inputs; ++i) {
                sum += inputs[i]*weights[o][i];
            }
            activation_function(sum, output_values[o], dout_per_dsum[o]);
        }
    }
    // must do forward step and initialize derr_per_dout before calling this
    void backward (flt learn_rate, flt derr_per_din[n_inputs]) {
        // zero the input (we'll sum inplace)
        //#pragma omp parallel_for
        for (int i = 0; i < n_inputs; ++i) {
            derr_per_din[i] = flt(0);
        }

        //float derr_per_dsum[n_inputs];
        // I should turn this loop inside out for parallelization
        for (int o = 0; o < n_outputs; ++o) {
            flt derr_per_dsum = derr_per_dout[o] * dout_per_dsum[o]; // d f(g(x)) / dx = f'(g(x)) * g'(x)

            //#pragma omp parallel for
            for (int i = 0; i < n_inputs; ++i) {
                derr_per_din[i] += weights[o][i] * derr_per_dsum; // dsum/din = weights[o][i]

                flt derr_per_weight = derr_per_dsum * input_values[i]; // dsum/dweight = input_values[i]

                weights[o][i] -= learn_rate * derr_per_weight ;
                //derr_per_weight_sums[o][i] += derr_per_weight ;
            }
        }
    }

    void backward_from_target(flt learn_rate, flt derr_per_din[n_inputs], const flt desired_outputs[n_inputs]){
        //#pragma omp parallel_for
        for(int o=0; o<n_outputs; ++o){
            derr_per_dout[o] = loss_function_der(desired_outputs[o], output_values[o]);
        }
        backward(learn_rate, derr_per_din);
    }

    // for Rprop
    /*
    void update_weights(){
        return;
        for (int o = 0; o < n_outputs; ++o)
            for (int i = 0; i < n_inputs; ++i){
                if((old_derr_per_weight_sums[o][i]>0) != (derr_per_weight_sums[o][i]>0)){// overshot, weight rollback
                    weight_step_size[o][i]*=step_shrink;
                    weights[o][i]=old_weights[o][i];
                }else{
                    old_weights[o][i]=weights[o][i];
                    weight_step_size[o][i]*=step_grow;
                    weight_step_size[o][i]=clamp(weight_step_size[o][i], step_min, step_max);

                    weights[o][i]-=derr_per_weight_sums[o][i]>0 ? weight_step_size[o][i] : -weight_step_size[o][i];

                    //weights[o][i] -= derr_per_weight_sums[o][i]*weight_step_size[o][i];

                    weights[o][i]=clamp(weights[o][i], flt(-4), flt(4));
                }
                old_derr_per_weight_sums[o][i]=derr_per_weight_sums[o][i];
                derr_per_weight_sums[o][i]=0;
            }
    }*/

    void print() {
        for (int o = 0; o < n_outputs; ++o) {
            for (int i = 0; i < n_inputs; ++i) {
                std::cout << weights[o][i] << "\t";
            }
            std::cout << std::endl;
        }
    }
};

layer<28*28, 600> input_layer;
layer<600,300> hidden_layer;
layer<300,10> output_layer;

void forward(const flt inputs[input_layer.n_inputs]) {
    input_layer.forward(inputs);
    hidden_layer.forward(input_layer.output_values);
    output_layer.forward(hidden_layer.output_values);
}
float loss(const flt desired_outputs[output_layer.n_outputs]) {
    flt error_sum{};
    for (int i = 0; i < output_layer.n_outputs; ++i) {
        error_sum += loss_function(desired_outputs[i], output_layer.output_values[i]);
    }
    return error_sum;
}

void backward(flt learn_rate, const flt desired_outputs[output_layer.n_outputs]) {
    output_layer.backward_from_target(learn_rate, hidden_layer.derr_per_dout, desired_outputs);
    hidden_layer.backward(learn_rate, input_layer.derr_per_dout);
    flt tmp[input_layer.n_inputs];
    input_layer.backward(learn_rate, tmp); // error deltas of the input are unused (todo: use them for visualization?)
}

/*
void update_weights(){
    for(auto &a: layers){
        a.update_weights();
    }
}*/

struct hand_written_digits_dataset{
    std::vector<flt> data;
    std::vector<uint8_t> labels;
    uint32_t data_width{}, data_height{}, count{};
    flt *get_image(size_t i){
        return data.data()+i*data_width*data_height;
    }
};

hand_written_digits_dataset train_set;
hand_written_digits_dataset test_set;

void load_data(hand_written_digits_dataset &result, const char *images_filename, const char *labels_filename){
    struct {
        uint32_t magic;
        uint32_t count;
        uint32_t width;
        uint32_t height;
    } images_header{};
    struct {
        uint32_t magic;
        uint32_t count;
    } labels_header{};
    std::cout<<"Reading images file "<<images_filename<<" and labels "<<labels_filename<<std::endl;
    std::ifstream f(images_filename, ios::binary|ios::in);
    f.read((char *)&images_header, sizeof(images_header));
    // convert
    images_header.magic=ntohl(images_header.magic);
    images_header.count=ntohl(images_header.count);
    result.data_width=images_header.width=ntohl(images_header.width);
    result.data_height=images_header.height=ntohl(images_header.height);

    std::ifstream fl(labels_filename, ios::binary|ios::in);
    fl.read((char *)&labels_header, sizeof(labels_header));
    labels_header.magic=ntohl(labels_header.magic);
    labels_header.count=ntohl(labels_header.count);

    if(labels_header.count!=images_header.count){
        std::cout<<"count mismatch "<<images_header.count<<" "<<labels_header.count<<std::endl;
        return;
    }
    std::cout<<"loading "<<images_header.count<<" images"<<std::endl;

    std::vector<uint8_t> images_data(images_header.count*images_header.width*images_header.height);
    f.read((char *)images_data.data(), images_data.size());
    if(!f.good()){
        std::cout<<"failed reading images"<<std::endl;
        return;
    }
    result.labels.resize(labels_header.count);
    fl.read((char *)result.labels.data(), result.labels.size());
    if(!fl.good()){
        std::cout<<"failed reading labels"<<std::endl;
        return;
    }
    result.data.resize(images_data.size());
    for(size_t i=0; i<images_data.size(); ++i){
        result.data[i]=images_data[i]*(1.0f/255.0f);
    }
    result.count=images_header.count;

    std::cout<<"read successful"<<std::endl;
}

int recognize(flt *inputs){
    forward(inputs);
    flt max_v=-1E10;
    int max_i=0;
    for(int i=0;i<10;++i){
    flt v=output_layer.output_values[i];
        if(v>max_v){
            max_v=v;
            max_i=i;
        }
    }
    return max_i;
}

void test(){
    flt loss_sum = 0;
    int errors = 0;
    for (int example = 0;  example < test_set.count; ++example) {
        errors+=recognize(test_set.get_image(example))!=test_set.labels[example];
    }
    std::cout<<"Error rate in test dataset: "<<(((float)errors)/test_set.count)<<std::endl;
}

void train() {
    for (int i = 0; i <= 1000000; ++i){
        flt loss_sum = 0;
        bool print = true;//!(i % 100);
        int errors = 0;
        for (int example = 0;  example < train_set.count; ++example) {
            flt *inputs=train_set.get_image(example);

            flt desired_outputs[10]{};
            int label=train_set.labels[example];
            if(label<0 || label>=10){
                std::cout<<"invalid label"<<std::endl;
            }
            desired_outputs[label]=1;

            forward(inputs);
            loss_sum += loss(desired_outputs);

            flt max_v=-1E10;
            int max_i=0;
            for(int i=0;i<10;++i){
                flt v=output_layer.output_values[i];
                if(v>max_v){
                    max_v=v;
                    max_i=i;
                }
            }
            errors+=max_i!=label;

            backward(0.01, desired_outputs);

            if(!(example % 1000))std::cout<<example<<std::endl;
        }// example loop
        if (print) {
            std::cout << "epoch " << i <<" loss sum over all examples:" << std::fixed << std::setprecision(6) << loss_sum <<" errors:" << errors << std::endl;
        }
        if (!errors) {
            std::cout << "Set fully recognized?!" << std::endl;
            break;
        }
        if(!(i%2)){
            test();
        }
    }
}

const char *symbols=" .,-+coC0#";

void print_image(flt *image, int w, int h){

    for(int i=0;i<h;++i){
        for(int j=0;j<w;++j){
            float a=*(image++);
            int av=clamp(int(a*10.0), 0, 9);
            std::cout<<symbols[av];
        }
        std::cout<<std::endl;
    }
}


int main()
{
    load_data(train_set, "../data/train-images.idx3-ubyte", "../data/train-labels.idx1-ubyte");
    load_data(test_set, "../data/t10k-images.idx3-ubyte", "../data/t10k-labels.idx1-ubyte");
    print_image(test_set.get_image(0), 28, 28);

    train();
    test();
    char c;
    std::cin >> c;
    return 0;
}