#include <stdio.h>
#include <vector>
#include <string>
#include <math.h>
#include "mnist_common.h"
float accuracy(
const array& predicted, 
const array& target)
 
{
    return 100 * count<float>(predicted == target) / target.
elements();
 
}
void naive_bayes_train(float *priors,
                       const array &train_feats,
 
                       const array &train_classes,
 
                       int num_classes)
{
    const int feat_len = train_feats.
dims(0);
 
    const int num_samples = train_classes.
elements();
 
    
    mu  = 
constant(0, feat_len, num_classes);
    sig2 = 
constant(0, feat_len, num_classes);
    for (int ii = 0; ii < num_classes; ii++) {
        
        sig2(
span,ii) = 
var(train_feats_ii, 0, 1) + 0.01;
        
        priors[ii] = (float)idx.
elements() / (float)num_samples;
    }
}
array naive_bayes_predict(
float *priors,
 
                          const array &test_feats, 
int num_classes)
 
{
    int num_test = test_feats.
dims(1);
 
    
    
    for (int ii = 0; ii < num_classes; ii++) {
        
        
        array Df = test_feats - Mu;
 
        
        log_probs(
span, ii) = 
log(priors[ii]) + 
sum(log_P).
T();
    }
    
    max(val, idx, log_probs, 1);
 
    return idx;
}
void benchmark_nb(
const array &train_feats, 
const array test_feats,
 
                  const array &train_labels, 
int num_classes)
 
{
    int iter = 25;
    float *priors = new float[num_classes];
    for (int i = 0; i < iter; i++) {
        naive_bayes_train(priors, mu, sig2, train_feats, train_labels, num_classes);
    }
    printf(
"Training time: %4.4lf s\n", 
timer::stop() / iter);
    for (int i = 0; i < iter; i++) {
        naive_bayes_predict(priors, mu, sig2, test_feats, num_classes);
    }
    printf(
"Prediction time: %4.4lf s\n", 
timer::stop() / iter);
    delete[] priors;
}
void naive_bayes_demo(bool console, int perc)
{
    array train_images, train_labels;
 
    array test_images, test_labels;
 
    int num_train, num_test, num_classes;
    
    float frac = (float)(perc) / 100.0;
    setup_mnist<false>(&num_classes, &num_train, &num_test,
                       train_images, test_images,
                       train_labels, test_labels, frac);
    int feature_length = train_images.
elements() / num_train;
 
    array train_feats = 
moddims(train_images, feature_length, num_train);
 
    array test_feats  = 
moddims(test_images , feature_length, num_test );
 
    
    float *priors = new float[num_classes];
    naive_bayes_train(priors, mu, sig2, train_feats, train_labels, num_classes);
    
    array res_labels = naive_bayes_predict(priors, mu, sig2, test_feats, num_classes);
 
    delete[] priors;
    
    printf("Trainng samples: %4d, Testing samples: %4d\n", num_train, num_test);
    printf("Accuracy on testing  data: %2.2f\n",
           accuracy(res_labels , test_labels));
    benchmark_nb(train_feats, test_feats, train_labels, num_classes);
    if (!console) {
        test_images = test_images.T();
        test_labels = test_labels.T();
        
        
    }
}
int main(int argc, char** argv)
{
    int device   = argc > 1 ? atoi(argv[1]) : 0;
    bool console = argc > 2 ? argv[2][0] == '-' : false;
    int perc     = argc > 3 ? atoi(argv[3]) : 60;
    try {
        naive_bayes_demo(console, perc);
        std::cerr << ae.
what() << std::endl;
    }
}