"""N64/N128 admission check: exact topology and independently rational interpolation.

Raw FP64 bytes are diagnostic: different LAPACK builds need not produce them
identically. The 2e-14 ceiling is the existing interpolation constraint policy,
not a universal LAPACK error theorem. This module never changes runtime weights.
"""
from fractions import Fraction as F
from functools import lru_cache

import numpy as np
from spacetime_gpu import table_hash

STRUCTURE_SHA = 'cd1842952c4dd4df82da8015454c25da9be12b983270236cc65b4044fa38e6e4'
STRUCTURES = {64: (STRUCTURE_SHA, 532),
              128: ('ee4ecb645674ae04eb87220a5217851405ccc0f45fbb0fe210499dcea5cd93b1', 1060)}
ATOL = 2e-14


def _transpose(a):
    return list(map(list, zip(*a)))


def _multiply(a, b):
    return [[sum(x*y for x,y in zip(row,column)) for column in zip(*b)] for row in a]


@lru_cache(None)
def exact_weights(kind, points):
    """Solve the small documented constrained fit using exact rational arithmetic."""
    if kind == 'face':
        a = [[F(1),F(s),F(t),F(s*t)] for t in (0,1) for s in (-1,0,1)]
        e = [[F(1),F(s),F(t),F(s)*F(t)] for s,t in points]
        b = [[sum(column)/4 for column in zip(*e)]]
        d = [[F(int(j==1)) for j in range(6)]]
    elif kind == 'corner':
        a = [[F(1),F(x),F(y),F(t)] for t in (0,1) for x,y in ((0,0),(-1,0),(1,0),(0,-1),(0,1))]
        e = [[F(1),F(x),F(y),F(t)] for x,y,t in points]
        b = [[sum(column)/4 for column in zip(*e)], [F(0),F(0),F(0),F(1)]]
        d = [[F(int(j==0)) for j in range(10)], [F(int(j==5)-int(j==0)) for j in range(10)]]
    else:
        raise ValueError('unsupported interpolation kind')
    kkt = [r+s for r,s in zip(_multiply(_transpose(a),a),_transpose(b))]
    kkt += [row+[F(0)]*len(b) for row in b]
    rhs = _transpose(a)+d
    n = len(kkt)
    augmented = [row+data for row,data in zip(kkt,rhs)]
    for j in range(n):
        pivot = next(i for i in range(j,n) if augmented[i][j])
        augmented[j],augmented[pivot] = augmented[pivot],augmented[j]
        divisor = augmented[j][j]
        augmented[j] = [v/divisor for v in augmented[j]]
        for i in range(n):
            if i != j:
                factor = augmented[i][j]
                augmented[i] = [x-factor*y for x,y in zip(augmented[i],augmented[j])]
    solution = [row[n:] for row in augmented]
    if _multiply(kkt,solution) != rhs:
        raise AssertionError('exact interpolation solve failed')
    return np.array(_multiply(e,solution[:4]), dtype=np.float64)


def validate_tables(plan, tables, *, allow_n128=False):
    """Reject topology, layout, padding, operator, or constraint changes separately."""
    structure = table_hash({key:value for key,value in tables.items() if key != 'weights'})
    admitted = (64,128) if allow_n128 is True else (64,)
    if plan.n not in admitted or structure != STRUCTURES[plan.n][0]:
        raise ValueError('frozen integer topology changed')
    w = tables['weights']
    if (w.dtype != np.dtype(np.float64) or w.shape != (STRUCTURES[plan.n][1],4,10) or len(plan.groups) != w.shape[0]
            or not all(a.flags.c_contiguous for a in tables.values()) or not np.isfinite(w).all()):
        raise ValueError('weight dtype/shape/layout/finiteness changed')
    errors = dict(coefficient=0., constant=0., mean=0., temporal=0.)
    for i,g in enumerate(plan.groups):
        size = len(g.samples)
        if size != tables['sizes'][i] or size != {'face':6,'corner':10}.get(g.kind) or len(g.points) != 4:
            raise ValueError('interpolation group coverage/size changed')
        if w[i,:,size:].tobytes() != np.zeros((4,10-size),np.float64).tobytes():
            raise ValueError('weight padding is not exact positive zero')
        actual = w[i,:,:size]
        errors['coefficient'] = max(errors['coefficient'],float(abs(actual-exact_weights(g.kind,g.points)).max()))
        errors['constant'] = max(errors['constant'],float(abs(actual.sum(axis=1)-1).max()))
        central = np.zeros(size); central[1 if g.kind=='face' else 0] = 1
        errors['mean'] = max(errors['mean'],float(abs(actual.mean(axis=0)-central).max()))
        if g.kind == 'corner':
            early,late = (g.points.index((0.,0.,t)) for t in (-.25,.25))
            slope = np.zeros(size); slope[0],slope[5] = -1,1
            errors['temporal'] = max(errors['temporal'],float(abs(2*(actual[late]-actual[early])-slope).max()))
    if any(error > ATOL for error in errors.values()):
        raise ValueError(f'exact interpolation/constraint check failed: {errors}; bound={ATOL}')
    return dict(status='passed', structure_sha256=structure, raw_table_sha256=table_hash(tables),
                coefficient_oracle='exact rational KKT, rounded once to FP64', atol=ATOL, rtol=0.,
                maximum_absolute_errors=errors, positive_zero_padding=True)