#!/usr/bin/env python3
"""Check the two worked counts using Python's standard library, without the site solver.

Usage: python3 verify-core-count-certificate.py core-count-certificate.json

For K=W+gamma*s*s^T, set F(a)=R(a)+lambda*(||K*a-s||^2-||s||^2).
Every predictor satisfying the training cap has R(a)>=F(a). If H is the
positive quadratic coefficient matrix of F and h is half its gradient at a0,
completing the square gives min F=F(a0)-h^T H^-1 h. Since H>=W>=mu I,
F(a0)-||h||^2/mu is a lower bound on every feasible predictor's risk.
Here a positive Gershgorin bound supplies mu without a spectral solver.

The mathematical bound is exact in real arithmetic. This verifier checks
the stored floating-point inputs with math.fsum and a 1e-10 comparison
reserve; it is not an interval-arithmetic certification. The margins from
the 0.1 target are much larger than that reserve. Neither these two counts
nor an adjacent pair establishes the first crossing of a nonmonotone curve.
"""
import json
import math
import sys
from pathlib import Path


def dot(a, b):
    return math.fsum(x*y for x, y in zip(a, b))


def verify(data):
    gamma = data['identity']['gamma']
    target = data['identity']['target']
    assert data['trainingCapRelativeToZero'] == 1
    reserve = 1e-10
    reports = []
    for record in data['records']:
        n = record['n']
        s, a = data['signal'][:n], record['alpha']
        factor = data['nuisanceFactor'][:n]
        assert len(a) == n and len(s) == n
        assert all(len(row) == i+1 and row[i] > 0 for i, row in enumerate(factor))
        w = [[dot(factor[i], factor[j]) for j in range(n)] for i in range(n)]
        k = [[w[i][j]+gamma*s[i]*s[j] for j in range(n)] for i in range(n)]
        wa = [dot(row, a) for row in w]
        error = [dot(row, a)-s[i] for i, row in enumerate(k)]
        survival = gamma*dot(s, a)
        risk = (survival-1)**2+dot(a, wa)
        energy = dot(s, s)
        train = dot(error, error)/energy
        assert abs(risk-record['risk']) < reserve
        assert abs(train-record['trainingRelativeToZero']) < reserve
        report = {'n': n, 'risk': risk, 'trainingRelativeToZero': train}
        if record['role'] == 'excluded':
            multiplier = record['multiplier']
            assert multiplier >= 0
            mu = min(w[i][i]-math.fsum(abs(x) for j, x in enumerate(w[i]) if j != i) for i in range(n))
            assert mu > 0
            half_gradient = [wa[i]+gamma*s[i]*(survival-1)+multiplier*dot(k[i], error) for i in range(n)]
            correction = dot(half_gradient, half_gradient)/mu
            lower = risk+multiplier*(dot(error, error)-energy)-correction
            assert lower-reserve > target
            report.update(lowerBound=lower, stationarityCorrection=correction,
                          eigenvalueLowerBound=mu, conclusion='target excluded under the training cap')
        else:
            assert record['role'] == 'feasible'
            assert train+reserve < 1 and risk+reserve < target
            report['conclusion'] = 'this predictor satisfies both requirements'
        reports.append(report)
    return {'identity': data['identity'], 'comparisonReserve': reserve,
            'certifiesFirstCrossing': False, 'checks': reports}


if __name__ == '__main__':
    source = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).with_name('core-count-certificate.json')
    print(json.dumps(verify(json.loads(source.read_text())), indent=2))