Export a full-precision inference bundle

"""Local Rice inference export, validation gate and digest-selected registry.

No external packages. This is a trusted local filesystem exercise, not MLflow.
"""
import hashlib
import json
import math
import re
from pathlib import Path
from vendor import baseline

HERE = Path(__file__).resolve().parent
CONTRACT = 'rice-measurements-v1'
POLICY = {'accuracy': [90, 100], 'class_recall': [85, 100]}


def encode(value):
    return (json.dumps(value, sort_keys=True, separators=(',', ':'),
                       allow_nan=False) + '\n').encode()


def digest(data):
    return hashlib.sha256(data).hexdigest()


def pairs(items):
    result = {}
    for key, value in items:
        if key in result:
            raise ValueError('duplicate JSON key')
        result[key] = value
    return result


def decode(data):
    return json.loads(data, object_pairs_hook=pairs,
                      parse_constant=lambda _: (_ for _ in ()).throw(ValueError('nonfinite JSON')))


def number(value):
    try:
        finite = type(value) in (float, int) and math.isfinite(value)
    except OverflowError:
        finite = False
    if not finite:
        raise ValueError('expected a finite non-Boolean number')
    return value


def validate(bundle):
    if not isinstance(bundle, dict):
        raise ValueError('expected a bundle object')
    if (type(bundle.get('format')) is not int or bundle.get('format') != 1 or bundle.get('contract') != CONTRACT
            or bundle.get('features') != baseline.FEATURES
            or bundle.get('labels') != baseline.LABELS
            or bundle.get('positive') != baseline.POSITIVE):
        raise ValueError('unsupported bundle contract')
    for field in ('weights', 'mean', 'std'):
        values = bundle.get(field)
        if not isinstance(values, list) or len(values) != len(baseline.FEATURES):
            raise ValueError('wrong parameter dimensions')
        for value in values:
            number(value)
    if any(value <= 0 for value in bundle['std']):
        raise ValueError('nonpositive standard deviation')
    number(bundle.get('bias'))
    return bundle


def predict(bundle, request, version):
    if not isinstance(request, dict) or set(request) != {'contract', 'features'}:
        raise ValueError('expected contract and features')
    features = request['features']
    if (request['contract'] != bundle['contract'] or not isinstance(features, dict)
            or set(features) != set(bundle['features'])):
        raise ValueError('feature contract mismatch')
    values = [number(features[name]) for name in bundle['features']]
    scaled = [(v - m) / s for v, m, s in zip(values, bundle['mean'], bundle['std'])]
    z = bundle['bias'] + sum(w * v for w, v in zip(bundle['weights'], scaled))
    if not all(math.isfinite(v) for v in scaled) or not math.isfinite(z):
        raise ValueError('nonfinite computation')
    probability = baseline.sigmoid(z)
    label = baseline.POSITIVE if z >= 0 else baseline.LABELS[1]
    return {'version': version, 'contract': CONTRACT, 'label': label,
            'positive_probability': probability}


def request_for(row):
    return {'contract': CONTRACT, 'features': dict(zip(baseline.FEATURES, row['x']))}


def export():
    parts = baseline.split_rows(baseline.load_rows())
    scaler = baseline.fit_scaler(parts['training'])
    model = baseline.train_logistic(parts['training'], scaler)
    bundle = {'format': 1, 'contract': CONTRACT, 'features': baseline.FEATURES,
              'labels': baseline.LABELS, 'positive': baseline.POSITIVE,
              'weights': model['weights'], 'bias': model['bias'], **scaler,
              'training': {'dataset_sha256': baseline.DATA_SHA256,
                           'split_salt': baseline.SALT, 'epochs': baseline.EPOCHS,
                           'learning_rate': baseline.LEARNING_RATE,
                           'trainer_sha256': digest(Path(baseline.__file__).read_bytes()),
                           'exporter_sha256': digest(Path(__file__).read_bytes())}}
    validate(bundle)
    # Match the source predictor on every validation row before retaining bytes.
    version = digest(encode(bundle))
    for row in parts['validation']:
        if predict(bundle, request_for(row), version)['label'] != baseline.predict_logistic(model, scaler, row):
            raise ValueError('export prediction mismatch')
    return encode(bundle)


def quality_gate(data, expected):
    if digest(data) != expected:
        raise ValueError('candidate digest mismatch')
    bundle = validate(decode(data))
    training = bundle.get('training', {})
    required = {'dataset_sha256': baseline.DATA_SHA256, 'split_salt': baseline.SALT,
                'epochs': baseline.EPOCHS, 'learning_rate': baseline.LEARNING_RATE,
                'trainer_sha256': digest(Path(baseline.__file__).read_bytes()),
                'exporter_sha256': digest(Path(__file__).read_bytes())}
    if training != required:
        raise ValueError('unsupported training identity')
    rows = baseline.split_rows(baseline.load_rows())['validation']
    matrix = {label: {other: 0 for other in baseline.LABELS} for label in baseline.LABELS}
    for row in rows:
        result = predict(bundle, request_for(row), expected)
        matrix[row['label']][result['label']] += 1
    correct = sum(matrix[label][label] for label in matrix)
    enough = lambda a, b, ratio: b > 0 and a * ratio[1] >= b * ratio[0]
    passed = enough(correct, len(rows), POLICY['accuracy']) and all(
        enough(matrix[label][label], sum(matrix[label].values()), POLICY['class_recall'])
        for label in matrix)
    return {'candidate_sha256': expected, 'partition': 'validation', 'policy': POLICY,
            'membership_sha256': digest(encode(sorted(row['key'] for row in rows))),
            'rows': len(rows), 'correct': correct, 'confusion': matrix, 'passed': passed}


def registered_file(root, version):
    if not isinstance(version, str) or not re.fullmatch('[0-9a-f]{64}', version):
        raise ValueError('select an explicit SHA-256 version')
    root = Path(root)
    if root.is_symlink():
        raise ValueError('symlink registry root')
    return root / (version + '.json')


def register(root, data, expected):
    record = quality_gate(data, expected)
    if not record['passed']:
        raise ValueError('candidate failed quality policy')
    path = registered_file(root, expected)
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.is_symlink():
        raise ValueError('symlink registration')
    try:
        with path.open('xb') as stream:
            stream.write(data)
    except FileExistsError:
        if path.read_bytes() != data:
            raise ValueError('conflicting or incomplete registration') from None
    return record


def load(root, expected):
    path = registered_file(root, expected)
    if path.is_symlink():
        raise ValueError('symlink registration')
    data = path.read_bytes()
    if digest(data) != expected:
        raise ValueError('registered bytes changed or incomplete')
    return validate(decode(data))


def batch(bundle, requests, version):
    if not isinstance(requests, list) or not 1 <= len(requests) <= 64:
        raise ValueError('batch size must be 1..64')
    # Return only a complete list. A bad row raises before any output is returned.
    return [predict(bundle, request, version) for request in requests]