// Author:  cirakg & claude ai
// Summary: Reed-Solomon error-correction codewords over GF(2^m) for QR, Data Matrix and Aztec, resumable for large symbols.
//
// A Field is GF(2^m) with a primitive polynomial and the exponent the
// generator's roots start at (QR: 0x11D, base 0; Data Matrix: 0x12D,
// base 1). Generator polynomials are cached per field and length. The
// error-correction symbols are the remainder of data * x^n divided by the
// generator: the standard systematic encoding, one multiply-accumulate per
// data symbol and EC symbol.
//
// Large Aztec symbols need GF(4096) and up to 832 EC symbols, far more than
// one event's watchdog budget, so every part can run in steps: the field's
// tables (grow), the generator (buildGenerator) and the division (Encoder).
// Fields of up to 256 elements are built at once (a few thousand
// bytecodes).

import Toybox.Lang;

module ReedSolomon {

    class Field {
        var size as Number;
        // exp has 2 * size entries once ready (exp[a + b] needs no modulo).
        var exp as Array<Number>;
        var log as Array<Number>;
        var base as Number;
        private var _primitive as Number;
        private var _filled as Number = 0;
        private var _x as Number = 1;
        private var _generators as Dictionary<Number, Array<Number>> = {} as Dictionary<Number, Array<Number>>;
        // The generator being built in steps: its length, coefficients so far
        // (leading 1 first) and the next root.
        private var _pendingN as Number = 0;
        private var _pending as Array<Number> = [] as Array<Number>;
        private var _pendingRoot as Number = 0;

        // primitive: the field polynomial including x^m (0x11D for GF(256)).
        function initialize(primitive as Number, fieldSize as Number, generatorBase as Number) {
            size = fieldSize;
            base = generatorBase;
            _primitive = primitive;
            exp = new [fieldSize - 1] as Array<Number>;
            log = new [fieldSize] as Array<Number>;
            log[0] = 0;
            if (fieldSize <= 256) {
                grow(fieldSize);
            }
        }

        function isReady() as Boolean {
            return exp.size() == 2 * size;
        }

        // Fill up to `count` more table entries; true when the field is ready.
        function grow(count as Number) as Boolean {
            if (isReady()) {
                return true;
            }
            var stop = _filled + count < size - 1 ? _filled + count : size - 1;
            var x = _x;
            for (var i = _filled; i < stop; i++) {
                exp[i] = x;
                log[x] = i;
                x <<= 1;
                if (x >= size) {
                    x ^= _primitive;
                }
            }
            _x = x;
            _filled = stop;
            if (_filled < size - 1) {
                return false;
            }
            // Doubled table: exp[size - 1 + j] = exp[j] for j = 0 .. size.
            var first = exp;
            exp = first.slice(0, null) as Array<Number>;
            exp.addAll(first);
            exp.add(first[0]);
            exp.add(first[1]);
            return true;
        }

        function multiply(a as Number, b as Number) as Number {
            if (a == 0 || b == 0) {
                return 0;
            }
            return exp[log[a] + log[b]];
        }

        // True when generator(n) is cached.
        function hasGenerator(n as Number) as Boolean {
            return _generators.hasKey(n);
        }

        // Build generator(n) with about `work` coefficient updates (at least
        // one root); true when it is cached.
        function buildGenerator(n as Number, work as Number) as Boolean {
            if (_generators.hasKey(n)) {
                return true;
            }
            if (_pendingN != n) {
                _pendingN = n;
                _pending = [1] as Array<Number>;
                _pendingRoot = 0;
            }
            var g = _pending;
            var budget = work;
            while (_pendingRoot < n && budget > 0) {
                // Multiply by (x - a^(base + i)); in GF(2^m) minus is plus.
                var logRoot = (base + _pendingRoot) % (size - 1);
                var next = new [g.size() + 1] as Array<Number>;
                next[0] = g[0];
                for (var k = 1; k < g.size(); k++) {
                    var c = g[k - 1];
                    next[k] = g[k] ^ (c == 0 ? 0 : exp[log[c] + logRoot]);
                }
                var last = g[g.size() - 1];
                next[g.size()] = last == 0 ? 0 : exp[log[last] + logRoot];
                budget -= g.size();
                g = next;
                _pendingRoot++;
            }
            _pending = g;
            if (_pendingRoot < n) {
                return false;
            }
            _generators[n] = g.slice(1, null) as Array<Number>;
            _pendingN = 0;
            _pending = [] as Array<Number>;
            return true;
        }

        // Coefficients of the generator with n roots, highest degree first,
        // without the leading 1.
        function generator(n as Number) as Array<Number> {
            buildGenerator(n, 0x7FFFFFFF);
            return _generators[n] as Array<Number>;
        }
    }

    // The n error-correction symbols of data (both as symbol values), in
    // one go. Jobs use Encoder to spread a long block over several steps.
    function ecc(field as Field, data as Array<Number>, n as Number) as Array<Number> {
        var encoder = new Encoder(field, data, n);
        encoder.step(data.size());
        return encoder.remainder;
    }

    // ecc() as a resumable computation: step() feeds a number of data
    // symbols; `remainder` holds the error-correction symbols at the end.
    // The field must be ready and generator(n) cached (or cheap).
    class Encoder {
        var remainder as Array<Number>;

        private var _exp as Array<Number>;
        private var _log as Array<Number>;
        private var _logG as Array<Number>;
        private var _data as Array<Number>;
        private var _i as Number = 0;

        function initialize(field as Field, data as Array<Number>, n as Number) {
            var g = field.generator(n);
            _exp = field.exp;
            _log = field.log;
            _data = data;
            remainder = new [n] as Array<Number>;
            _logG = new [n] as Array<Number>;
            for (var k = 0; k < n; k++) {
                remainder[k] = 0;
                _logG[k] = g[k] == 0 ? -1 : _log[g[k]];
            }
        }

        // Up to `symbols` more data symbols; true when all are in.
        function step(symbols as Number) as Boolean {
            var rem = remainder;
            var n = rem.size();
            var exp = _exp;
            var logG = _logG;
            var stop = _i + symbols < _data.size() ? _i + symbols : _data.size();
            for (var i = _i; i < stop; i++) {
                var factor = _data[i] ^ rem[0];
                // Shift the remainder left by one symbol (natively).
                rem = rem.slice(1, null) as Array<Number>;
                rem.add(0);
                if (factor != 0) {
                    var lf = _log[factor];
                    for (var k = 0; k < n; k++) {
                        var lg = logG[k];
                        if (lg >= 0) {
                            rem[k] = rem[k] ^ exp[lf + lg];
                        }
                    }
                }
            }
            remainder = rem;
            _i = stop;
            return _i >= _data.size();
        }
    }
}