/*
 * minimal_mnist_ntuple.cpp
 *
 * Alex Sverdlov; alex@theparticle.com
 * 
 * g++ -std=c++17 -O3 minimal_mnist_ntuple.cpp -o minimal_mnist_ntuple
 *
 * ./minimal_mnist_ntuple train=/home/alex/datasets/EMNIST.uncompressed/emnist-digits-train-images-idx3-ubyte train_lbl=/home/alex/datasets/EMNIST.uncompressed/emnist-digits-train-labels-idx1-ubyte test=/home/alex/datasets/EMNIST.uncompressed/emnist-digits-test-images-idx3-ubyte test_lbl=/home/alex/datasets/EMNIST.uncompressed/emnist-digits-test-labels-idx1-ubyte
 *
 * get the dataset here: http://yann.lecun.com/exdb/mnist/
 */

#include <iostream>
#include <fstream>
#include <cstdio>
#include <numeric>
#include <cmath>
#include <vector>
#include <cstring>
#include <map>
#include <type_traits>
#include <cstdlib>
#include <algorithm>
#include <random>
#include <memory>

// MNIST file reader.
struct MNISTReader {
    std::ifstream is;
    std::vector<int> dim;
    int siz;
    std::vector<unsigned char> buf;
    MNISTReader(const std::string& f) : is(f) {
        if(!is) return;
        unsigned char b[4];
        is.read((char*)b,4);
        if(!is || b[0] || b[1] || b[2] != 0x08 || b[3] > 3){ is.setstate(std::ios::failbit); return; }
        dim.resize(b[3]);
        for(auto p=dim.begin();p!=dim.end();p++){
            is.read((char*)b,4);  // big endian integers
            if(!is) return;
            union { unsigned char c[4]; int i; } t { {b[3],b[2],b[1],b[0]} };
            *p = t.i;
        }
        siz = std::accumulate(++dim.begin(),dim.end(),1, std::multiplies<>() );
        buf.resize(siz);
    }
    bool next(){
        if(!is) return false;
        is.read((char*)buf.data(),siz);
        if(!is) return false;
        return true;
    }
};

/* 
 * M is num of modules
 * N is num of bits per module
 * L is num of classes
 * idx is set of indexes, of size M * N
 * tbl is table, of size: M * (1<<N) * L
 * buf is the buffer we wish to classify.
 */
int classify(int M, int N, int L, std::vector<int>& idx, std::vector<int>& tbl, std::vector<unsigned char>& buf){
  std::vector<int> agg(10,0);
  for(int m=0;m<M;m++){     // for each module
    int addr = 0;
    for(int n=0;n<N;n++)    // for each index
      addr += (buf[ idx[m*N+n] ]>=128 ? 1:0) << n;  // calculate address into m-module table.
    for(int l=0;l<L;l++)    // for each target label
      agg[l] += tbl[ m * (1<<N)*L + addr*L + l ];
  }
  return std::max_element(agg.begin(),agg.end()) - agg.begin();
}

/*
 * training 
 */
int train(int M, int N, int L, std::vector<int>& idx, std::vector<int>& tbl, std::vector<unsigned char>& buf, int l){
  for(int m=0;m<M;m++){     // for each module
    int addr = 0;
    for(int n=0;n<N;n++)    // for each index
      addr += (buf[ idx[m*N+n] ]>=128 ? 1:0) << n;  // calculate address into m-module table.
    tbl[ m*(1<<N)*L + addr*L + l ]++;
  }
}

int main(int argc,char** argv){
    srand(42);

    // read command line args
    std::map<std::string,std::string> args;
    for(int i=0;argv[i];i++){
        char* eq = strchr(argv[i],'=');
        if(eq){
            *eq = '\0'; eq++;
            args[ argv[i] ] = eq;
        }else{
            args[ argv[i] ] = "1";
        }
    }

    // load training and test data
    MNISTReader v{args["train"]}, l{args["train_lbl"]};
    MNISTReader vt{args["test"]}, lt{args["test_lbl"]};

    int M = 500;        // M-modules
    int N = 10;         // N-bits each
    int L = 10;         // MNIST digits
    std::vector<int> idx(M*N,0);                // index sets
    std::vector<int> tbl(M * (1<<N) * L,0);     // n-table tables

    // randomly select M*N indexes
    // we could be smarter about this; if we pick only indexes that have 
    // non-zero values in training data accuracy will increase.
    for(int i=0;i<M*N;i++)
        idx[i] = rand()%v.siz;

    int rows = v.dim[1], cols = v.dim[2];

    // training loop
    while(v.next() && l.next()){
        int o = classify(M,N,L,idx,tbl,v.buf);
        if(o != l.buf[0])
            train(M,N,L,idx,tbl,v.buf,l.buf[0]);
    }
    
    // set confusion matrix
    std::vector<int> confmat(L*L,0);

    int cnt = 0;
    // testing loop
    while(vt.next() && lt.next()){
        int o = classify(M,N,L,idx,tbl,vt.buf);
        confmat[ lt.buf[0]*L + o]++;
        cnt++;

        if(o != lt.buf[0] && cnt % 100 == 0){   // display some wrong ones.
          for(int i=0;i<cols;i++){
            for(int j=0;j<rows;j++)
              printf("%d",vt.buf[j*cols+i] > 0 ? 1:0);
            printf("\n");
          }
          printf("\n");
          printf("predicted: %d, actual: %d, dim:%dx%d = %d\n",o, lt.buf[0], rows,cols, v.siz);
        }
    }
    
    printf("CONFUSION MATRIX (rows actual, columns predicted).\n");
    int tot=0, mtch=0;
    for(int i=0;i<L;i++){
      for(int j=0;j<L;j++){
        tot += confmat[i*L+j];
        if(i==j) mtch+=confmat[i*L+j];
        printf("%7d",confmat[i*L+j]);
      }
      printf("\n");
    }
    printf("\n");
    printf("accuracy: %f\n",1.0*mtch/tot);
}


/*

0000000000000000000000000000
0000000000000000000001110000
0000000000000000000011111000
0000000000000000001111111000
0000000000000000001111111000
0000000000000000011111000000
0000000000000000011111000000
0000000000001111111110000000
0000000001111111111110000000
0000000011111111111110000000
0000000111111111111111000000
0000001111111011111111100000
0000001111110011111111110000
0000111111100111110011110000
0000111111000111110011110000
0001111100001111100111110000
0001111100011111111111100000
0001111100011111111111100000
0000111111111111111111000000
0000111111111111111110000000
0000011111111111111000000000
0000001111111111100000000000
0000000011111100000000000000
0000000011111100000000000000
0000000011111000000000000000
0000000011110000000000000000
0000000011100000000000000000
0000000000000000000000000000

predicted: 4, actual: 0, dim:28x28 = 784
...

0000000000000000000000000000
0000000000110000000000000000
0000000001111100000000000000
0000000011111110000000000000
0000000111111110000000000000
0000001111111111000111100000
0000001111001111000111110000
0000001111001110000111110000
0000011110000100001111100000
0000011110000000011111100000
0000011110000000011111100000
0000011110000000111111100000
0000011110000001111111100000
0000001111000011111111000000
0000001111000111111111000000
0000001111111111111111000000
0000001111111111111111000000
0000001111111111011111000000
0000000111111100111110000000
0000000011100000111110000000
0000000000000000111110000000
0000000000000000111110000000
0000000000000000111100000000
0000000000000000111100000000
0000000000000000111110000000
0000000000000000111100000000
0000000000000000011000000000
0000000000000000000000000000

predicted: 4, actual: 9, dim:28x28 = 784


...
0000000000000000000000000000
0000000000000000000111000000
0000000000000000001111000000
0000000000000000111111000000
0000000000000011111111000000
0000000000000111111111000000
0000000000011111111110000000
0000000000111111111110000000
0000000001111111001110000000
0000000011111100111110000000
0000000111111101111110000000
0000000111110111111110000000
0000001111111111111110000000
0000011111111111111110000000
0000011111111111111100000000
0000001111111101111100000000
0000000111100001111100000000
0000000000000011111000000000
0000000000000011111000000000
0000000000000111110000000000
0000000000000111110000000000
0000000000000111100000000000
0000000000001111100000000000
0000000000001111100000000000
0000000000001111000000000000
0000000000011111000000000000
0000000000001110000000000000
0000000000000000000000000000

predicted: 4, actual: 9, dim:28x28 = 784

...

CONFUSION MATRIX (rows actual, columns predicted).
   3797      2     25     15     34     65      9      1     43      9
      0   3775      9     24     10     52      0     30     98      2
     18      2   3743     59     55     25      7     17     66      8
     13      0     26   3792      5     56      0     16     71     21
     11      5     36      3   3774      7     22      7     60     75
     12      5     15    109     39   3698     17      6     77     22
     27     15     84      1     21     56   3754      0     42      0
      3      1      7     21     16      7      0   3809     51     85
      4      0     20     36     17     49      3      6   3835     30
      5      1      9     39     47     10      0     67     60   3762

accuracy: 0.943475
*/