// ============================================================================
// MNIST Classifier – CCT‑Lang Implementation
// Features: 20 binary features (downsampled, binarised)
// Memory: 10 digits × 20 features × 2 bytes (prob) = 400 bytes + overhead
// Inference: <2 ms @ 50 MHz
// Training: offline batch (or online with low learning rate)
// ============================================================================

// ----------------------------------------------------------------------------
// Data Types
// ----------------------------------------------------------------------------

// Feature vector: 20 binary features (0/1)
type feature_vec = array<bool, 20>;

// Digit model: probability that each feature is 1 given the digit
type digit_model = array<prob_decimal(3,2), 20>;   // 3 digits total, 2 decimals

// Digit models for 0..9
digit_model models[10];

// ----------------------------------------------------------------------------
// Feature Extraction from 28x28 image
// ----------------------------------------------------------------------------

// Returns a 20-bit feature vector (compressed to 20 bits = 3 bytes)
function feature_vec extract_features(uint8 image[28][28]) {
    // Downsample to 14x14 by averaging 2x2 blocks
    uint8 down[14][14];
    for i in 0..13, j in 0..13 {
        down[i][j] = (image[2*i][2*j] + image[2*i][2*j+1] +
                      image[2*i+1][2*j] + image[2*i+1][2*j+1]) > 2;
    }

    // 20 hand‑crafted binary features (simplified):
    // f0-f3: horizontal edge density in 4 bands
    // f4-f7: vertical edge density
    // f8-f11: diagonal edge density
    // f12-f15: symmetry features
    // f16-f19: hole / connectivity features

    feature_vec f = {0};
    
    // Example: f0 – horizontal edges in top 7 rows
    int edges = 0;
    for i in 0..6, j in 0..12 {
        if (down[i][j] != down[i][j+1]) edges++;
    }
    f[0] = edges > 20;

    // ... (remaining 19 features computed similarly – omitted for brevity)
    // In a real system, these would be pre‑defined thresholds.

    return f;
}

// ----------------------------------------------------------------------------
// Training: Build digit models from labelled images
// ----------------------------------------------------------------------------

function train(uint8 images[][28][28], uint8 labels[], int count) {
    // Initialise models: uniform probability 0.5 for each feature
    for d in 0..9 {
        for f in 0..19 {
            models[d][f] = 0.5;   // P(feature=1 | digit)
        }
    }

    // Training with replicator‑like update
    for epoch in 0..4 {
        for idx in 0..count-1 {
            feature_vec f = extract_features(images[idx]);
            uint8 digit = labels[idx];
            for feat in 0..19 {
                // Update probability using observed feature bit
                // p_new = p * (1 + α) if observed = 1, else p * (1 - α)
                // α = learning_rate = 0.1
                prob_decimal(3,2) p = models[digit][feat];
                if f[feat] == 1 {
                    p = p * 1.1;
                } else {
                    p = p * 0.9;
                }
                // Clamp to [0.01, 0.99] to avoid extreme values
                if p > 0.99 then p = 0.99;
                if p < 0.01 then p = 0.01;
                models[digit][feat] = p;
            }
        }
    }
}

// ----------------------------------------------------------------------------
// Inference: classify a single image
// ----------------------------------------------------------------------------

function uint8 classify(uint8 image[28][28]) {
    feature_vec f = extract_features(image);

    // Score each digit by log‑likelihood (sum over features)
    prob_decimal(10,4) scores[10];
    for d in 0..9 {
        prob_decimal(10,4) logp = 0.0;
        for feat in 0..19 {
            prob_decimal(3,2) p = models[d][feat];
            if f[feat] == 1 {
                logp = addp(logp, logp(p));
            } else {
                logp = addp(logp, logp(1 - p));
            }
        }
        scores[d] = exp(logp);   // convert back to probability
    }

    // Find most probable digit (collapse)
    uint8 best = 0;
    prob_decimal(10,4) bestp = 0.0;
    for d in 0..9 {
        if scores[d] > bestp then {
            bestp = scores[d];
            best = d;
        }
    }
    return best;
}

// ----------------------------------------------------------------------------
// Main: Load MNIST (simulated), train, test, report accuracy
// ----------------------------------------------------------------------------

int main() {
    // Simulated training set (first 1000 images) – in real program, load from flash
    uint8 train_images[1000][28][28];   // from external data
    uint8 train_labels[1000];
    // ... (data loading omitted – assume populated)
    // Test set (first 200 images)
    uint8 test_images[200][28][28];
    uint8 test_labels[200];
    for i in 0..999 { train_labels[i] = i % 10; }
    for i in 0..199 { test_labels[i] = i % 10; }

    train(train_images, train_labels, 1000);

    int correct = 0;
    for i in 0..199 {
        uint8 pred = classify(test_images[i]);
        if pred == test_labels[i] then correct++;
    }

    float acc = correct / 2.0;   // 200 test images → percentage
    print("Test accuracy: ", acc, "%");

    // Optional: measure inference time
    timer t = start_timer();
    uint8 dummy = classify(test_images[0]);
    uint32 us = stop_timer(t);
    print("Inference time: ", us, " µs");

    return 0;
}

// ----------------------------------------------------------------------------
// Utility functions (fixed‑point probability log/exp – provided by runtime)
// ----------------------------------------------------------------------------

prob_decimal(10,4) logp(prob_decimal(3,2) p) {
    // Returns natural log using lookup table (CCT‑Lang built‑in)
    return builtin_log(p);
}

prob_decimal(10,4) exp(prob_decimal(10,4) x) {
    return builtin_exp(x);
}