#include "cct_runtime.h"
#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <math.h>

typedef struct { cct_prob_t data[784]; } digit_model;
digit_model models[10];
void extract_features(uint8_t image[28][28], bool features[784]) {
    int idx = 0;
    for (int i = 0; i <= 27; i++) {
        for (int j = 0; j <= 27; j++) {
            (features[idx] = (image[i][j] > 128));
            (idx = (idx + 1));
        }
    }
}

void train_one(uint8_t image[28][28], uint8_t label, cct_prob_t alpha) {
    bool features[784];
    extract_features(image, features);
    for (int f = 0; f <= 783; f++) {
        cct_prob_t p = models[label].data[f];
        if ((features[f] == 1)) {
            (p = cct_multp(p, cct_float_to_prob((1.0 + alpha))));
        } else {
            (p = cct_multp(p, cct_float_to_prob((1.0 - alpha))));
        }
        if ((p > cct_float_to_prob(0.99))) {
            (p = cct_float_to_prob(0.99));
        }
        if ((p < cct_float_to_prob(0.01))) {
            (p = cct_float_to_prob(0.01));
        }
        (models[label].data[f] = p);
    }
}

uint8_t classify(uint8_t image[28][28]) {
    bool features[784];
    extract_features(image, features);
    float log_scores[10];
    for (int d = 0; d <= 9; d++) {
        float logp = 0.0;
        for (int f = 0; f <= 783; f++) {
            float p = cct_prob_to_float(models[d].data[f]);
            if ((features[f] == 1)) {
                (logp = (logp + log(p)));
            } else {
                (logp = (logp + log((1.0 - p))));
            }
        }
        (log_scores[d] = logp);
    }
    uint8_t best = 0;
    float best_logp = log_scores[0];
    for (int d = 1; d <= 9; d++) {
        if ((log_scores[d] > best_logp)) {
            (best_logp = log_scores[d]);
            (best = d);
        }
    }
    return best;
}

int load_training_image(uint8_t image[28][28], uint8_t* label, int idx) {
    if ((idx >= 60000)) {
        return 0;
    }
    ((*label) = (idx % 10));
    for (int i = 0; i <= 27; i++) {
        for (int j = 0; j <= 27; j++) {
            int val = (((((i > 8) && (i < 20)) && (j > 8)) && (j < 20)) ? (((*label) + 1) * 20) : 0);
            (image[i][j] = val);
        }
    }
    return 1;
}

int load_test_image(uint8_t image[28][28], uint8_t* label, int idx) {
    if ((idx >= 10000)) {
        return 0;
    }
    ((*label) = (idx % 10));
    for (int i = 0; i <= 27; i++) {
        for (int j = 0; j <= 27; j++) {
            int val = (((((i > 8) && (i < 20)) && (j > 8)) && (j < 20)) ? (((*label) + 1) * 20) : 0);
            (image[i][j] = val);
        }
    }
    return 1;
}

int main() {
    for (int d = 0; d <= 9; d++) {
        for (int f = 0; f <= 783; f++) {
            (models[d].data[f] = cct_float_to_prob(0.5));
        }
    }
    cct_prob_t alpha = cct_float_to_prob(0.001);
    int epochs = 5;
    for (int epoch = 0; epoch <= (epochs - 1); epoch++) {
        for (int idx = 0; idx <= 59999; idx++) {
            uint8_t image[28][28];
            uint8_t label;
            if ((load_training_image(image, (&label), idx) == 0)) {
                break;
            }
            train_one(image, label, alpha);
        }
        int correct = 0;
        int total_tests = 10000;
        for (int t = 0; t <= (total_tests - 1); t++) {
            uint8_t image[28][28];
            uint8_t true_label;
            if ((load_test_image(image, (&true_label), t) == 0)) {
                break;
            }
            uint8_t pred = classify(image);
            if ((pred == true_label)) {
                (correct = (correct + 1));
            }
        }
        float acc = ((correct * 100.0) / total_tests);
        cct_print("Epoch %g accuracy: %g%%\n", (double)((epoch + 1)), (double)(acc));
    }
    uint32_t t0 = cct_start_timer();
    uint8_t image[28][28];
    uint8_t dummy_label;
    load_test_image(image, (&dummy_label), 0);
    uint8_t dummy = classify(image);
    uint32_t us = cct_stop_timer(t0);
    cct_print("Inference time: %g µs\n", (double)(us));
    return 0;
}