import json
import numpy as np

# Deterministic numerical check for the proposed weak residual and entropy term
# on stationary Burgers shock profiles.  This is intentionally a toy mechanism
# test rather than a claim of trained-PINN performance.
T = 1.0
x = np.linspace(-1.0, 1.0, 2001)
t = np.linspace(0.0, T, 161)
X = x[None, :]
TT = t[:, None]

def trap2(a):
    return np.trapz(np.trapz(a, x, axis=1), t, axis=0)

def profile(eps, kind):
    # Correct entropy shock: +1 on x<0 and -1 on x>0.
    # Anti-entropic orientation: -1 on x<0 and +1 on x>0.
    sign = 1.0 if kind == 'entropy_shock' else -1.0
    return -sign * np.tanh(X / eps)

def test_fields(center, width):
    z = np.exp(-((X - center) / width) ** 2)
    a = TT * (T - TT)  # nonnegative and zero at both temporal endpoints
    phi_t = (T - 2.0 * TT) * z
    phi_x = a * z * (-2.0 * (X - center) / width**2)
    return phi_t, phi_x

def evaluate(eps, kind, width=0.24, center=0.0):
    u = profile(eps, kind)
    ux = -(1.0 / eps) / np.cosh(X / eps) ** 2
    strong_rms = np.sqrt(np.mean((u * ux) ** 2))
    phi_t, phi_x = test_fields(center, width)
    f = 0.5 * u**2
    eta = 0.5 * u**2
    q = u**3 / 3.0
    weak_r = -trap2(u * phi_t + f * phi_x)
    entropy_s = trap2(eta * phi_t + q * phi_x)
    return float(strong_rms), float(weak_r), float(entropy_s)

def log_slope(xs, ys):
    return float(np.polyfit(np.log(xs), np.log(np.maximum(np.abs(ys), 1e-15)), 1)[0])

def main():
    epses = np.array([0.20, 0.12, 0.08, 0.05, 0.03, 0.02])
    rows = []
    strong, weak = [], []
    shock_s, anti_s = [], []
    for eps in epses:
        a, b, c = evaluate(eps, 'entropy_shock')
        d, e, f = evaluate(eps, 'anti_entropic')
        strong.append(a); weak.append(abs(b)); shock_s.append(c); anti_s.append(f)
        rows.append({'epsilon': float(eps), 'strong_rms': a, 'weak_abs': abs(b),
                     'S_entropy_shock': c, 'S_anti_entropic': f})

    # Quantitative predictions from the mechanism:
    # (1) strong RMS ~ eps^-1/2; (2) weak residual ~ 0; (3) for the
    # distributional inequality d_t eta+d_x q <= 0, integrated S must be >= 0.
    # Therefore relu(S)^2 as written penalizes the physical entropy shock.
    widths = [0.12, 0.24, 0.40]
    width_rows = []
    for width in widths:
        _, _, a = evaluate(0.04, 'entropy_shock', width=width)
        _, _, b = evaluate(0.04, 'anti_entropic', width=width)
        width_rows.append({'test_width': width, 'S_entropy_shock': a,
                           'S_anti_entropic': b})

    out = {
        'predictions': {
            'strong_RMS_vs_epsilon': {'predicted_log_slope': -0.5,
                                     'observed_log_slope': log_slope(epses, strong),
                                     'relative_error': abs(log_slope(epses, strong) + 0.5) / 0.5},
            'weak_RMS_vs_epsilon': {'predicted': 'near zero for Rankine-Hugoniot layer',
                                    'observed_max_abs': max(weak)},
            'entropy_integral_sign': {
                'mathematical_prediction': 'S >= 0 for entropy-admissible solutions under the displayed convention',
                'observed_shock_sign': 'positive', 'observed_anti_entropic_sign': 'negative',
                'penalty_relu_S_selects_physical_shock': False}
        },
        'sweep': rows,
        'test_width_sweep': width_rows,
        'worked_mechanism': False
    }
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()