import json, random
from dataclasses import dataclass
import numpy as np

@dataclass
class LayerReport:
    lower: np.ndarray
    upper: np.ndarray
    active: int
    inactive: int
    unstable: int
    retained: int
    pruned_contribution: int


def affine_interval(W, b, lo, hi):
    W, b = np.asarray(W, float), np.asarray(b, float)
    lo, hi = np.asarray(lo, float), np.asarray(hi, float)
    lower = np.maximum(W, 0) @ lo + np.minimum(W, 0) @ hi + b
    upper = np.maximum(W, 0) @ hi + np.minimum(W, 0) @ lo + b
    return lower, upper


def interval_propagate(hidden_weights, hidden_biases, xlo, xhi):
    """Interval bounds for hidden preactivations and post-ReLU variables."""
    lo, hi = np.asarray(xlo, float), np.asarray(xhi, float)
    out = []
    for W, b in zip(hidden_weights, hidden_biases):
        L, U = affine_interval(W, b, lo, hi)
        zlo, zhi = np.maximum(0, L), np.maximum(0, U)
        out.append((L, U, zlo, zhi))
        lo, hi = zlo, zhi
    return out


def prune_report(hidden_weights, hidden_biases, xlo, xhi, tau=0.0):
    """Classify hidden neurons; only unstable neurons need SDP activation variables."""
    bounds = interval_propagate(hidden_weights, hidden_biases, xlo, xhi)
    reports = []
    for l, (L, U, _, _) in enumerate(bounds):
        # Make categories disjoint at the degenerate boundary.
        active_mask = L >= 0
        inactive_mask = (~active_mask) & (U <= 0)
        unstable_mask = ~(active_mask | inactive_mask)
        if l + 1 < len(hidden_weights):
            outgoing = np.max(np.abs(hidden_weights[l + 1]), axis=0)
        else:
            outgoing = np.ones(len(L))
        magnitude = np.maximum(np.abs(L), np.abs(U))
        contribution = outgoing * magnitude
        keep = unstable_mask & (contribution > tau)
        reports.append(LayerReport(
            L, U, int(active_mask.sum()), int(inactive_mask.sum()),
            int(unstable_mask.sum()), int(keep.sum()),
            int((unstable_mask & ~keep).sum())))
    return bounds, reports


def forward(weights, biases, x):
    z = np.asarray(x, float)
    for k, (W, b) in enumerate(zip(weights, biases)):
        z = np.asarray(W) @ z + np.asarray(b)
        if k < len(weights) - 1:
            z = np.maximum(0, z)
    return z


def output_interval(weights, biases, xlo, xhi):
    hidden_bounds = interval_propagate(weights[:-1], biases[:-1], xlo, xhi)
    L, U = hidden_bounds[-1][2], hidden_bounds[-1][3]
    return affine_interval(weights[-1], biases[-1], L, U)


def make_net(seed=7):
    rng = np.random.default_rng(seed)
    return ([rng.normal(0, 1.0, (12, 2)), rng.normal(0, .8, (10, 12)), rng.normal(0, .7, (2, 10))],
            [rng.normal(0, .9, 12), rng.normal(0, .7, 10), rng.normal(0, .3, 2)])


def main():
    random.seed(11); np.random.seed(11)
    W, b = make_net(); x = np.array([.15, -.2]); eps = .3
    xlo, xhi = x - eps, x + eps
    hiddenW, hiddenb = W[:-1], b[:-1]
    print('CORE_CHECK')
    for tau in [0.0, 1e-6, 1e-2, 1e-1, 3e-1, 1.0]:
        _, reps = prune_report(hiddenW, hiddenb, xlo, xhi, tau)
        print(json.dumps({'tau': tau, 'unstable': sum(r.unstable for r in reps),
            'fixed_sign': sum(r.active + r.inactive for r in reps),
            'retained': sum(r.retained for r in reps),
            'contribution_pruned': sum(r.pruned_contribution for r in reps),
            'sdp_variable_proxy': 2 * sum(r.retained for r in reps)}))
    bounds, reps = prune_report(hiddenW, hiddenb, xlo, xhi, 0.0)
    max_fixed_err = 0.0; samples = []
    for xx in np.linspace(xlo[0], xhi[0], 81):
        for yy in np.linspace(xlo[1], xhi[1], 81):
            v = np.array([xx, yy]); z = v
            for l, (Wl, bl) in enumerate(zip(hiddenW, hiddenb)):
                a = Wl @ z + bl; L, U = bounds[l][0], bounds[l][1]
                pred = np.where(L >= 0, a, np.where(U <= 0, 0, np.maximum(0, a)))
                max_fixed_err = max(max_fixed_err, float(np.max(np.abs(pred - np.maximum(0, a)))))
                z = np.maximum(0, a)
            samples.append(forward(W, b, v))
    samples = np.asarray(samples); outL, outU = output_interval(W, b, xlo, xhi)
    violations = int(np.sum((samples < outL - 1e-9) | (samples > outU + 1e-9)))
    print('SOUNDNESS', json.dumps({'grid_points': len(samples), 'max_fixed_substitution_error': max_fixed_err,
        'output_interval_lower': outL.tolist(), 'output_interval_upper': outU.tolist(),
        'sample_interval_violations': violations}))
    print('SUMMARY', json.dumps({'baseline_hidden_sdp_proxy': 2 * sum(len(v) for v in hiddenb),
        'exact_pruned_hidden_proxy': 2 * sum(r.retained for r in reps),
        'fixed_sign_neurons': sum(r.active + r.inactive for r in reps)}))

if __name__ == '__main__': main()