"""Inference capacity, batching, and autoscaling: worked experiments and reference solutions. CPU checks."""

# 1. Define and check the inference workload
# Step 1 — 1. Define and check the inference workload: The timer includes dispatch and completion for already resident...
import math
import time
import numpy as np
import jax
import jax.numpy as jnp
# Draw pseudorandom samples for `rng` using the explicit RNG state.
rng=np.random.default_rng(82)
# Create device-backed JAX array `W1`.
W1=jnp.asarray(rng.normal(0,.1,(64,32)).astype(np.float32))
# Create device-backed JAX array `W2`.
W2=jnp.asarray(rng.normal(0,.1,(32,8)).astype(np.float32))
# Define and JIT-compile `infer(inputs)` so XLA traces and fuses the operations:
@jax.jit
def infer(inputs):
    return jnp.tanh(inputs@W1)@W2
def measure(batch_size,repeats=40):
    host=rng.normal(size=(batch_size,64)).astype(np.float32)
    # Create device-backed JAX array `device`.
    device=jnp.asarray(host)
    began=time.perf_counter()
    # Synchronize host execution until asynchronous device computation completes.
    first=infer(device).block_until_ready()
    first_ms=(time.perf_counter()-began)*1000
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected=np.tanh(host@np.asarray(W1))@np.asarray(W2)
    np.testing.assert_allclose(first,expected,rtol=2e-5,atol=2e-6)
    samples=[]
    for _ in range(repeats):
        began=time.perf_counter()
        # Synchronize host execution until asynchronous device computation completes.
        infer(device).block_until_ready()
        samples.append((time.perf_counter()-began)*1000)
    p50=float(np.percentile(samples,50))
    return {'batch':batch_size,'first_call_ms':first_ms,'samples_ms':samples,
            'p50_ms':p50,'p95_ms':float(np.percentile(samples,95)),
            'steady_examples_per_second':1000*batch_size/p50}

# 2. Measure each batch shape separately
# Step 2 — 2. Measure each batch shape separately: Do not assume a larger batch is faster per request.
measurements=[measure(batch) for batch in (1,8,32)]
for report in measurements:
    assert len(report['samples_ms'])==40
    # Ensure all array elements remain finite: `all(value>0 and np.isfinite(value) for value in report['samples_m...`
    assert all(value>0 and np.isfinite(value) for value in report['samples_ms'])
    print({key:value for key,value in report.items() if key!='samples_ms'})
service_ms=measurements[0]['p50_ms']

# 3. Model a queue with an independent hand-check
# Step 3 — 3. Model a queue with an independent hand-check: This simulator assumes one serial worker and a constant service...
def simulate_queue(arrival_ms,service_ms):
    # Convert `arrivals` to a host NumPy array for inspection or verification.
    arrivals=np.asarray(arrival_ms,dtype=float)
    # Guard input contract (`arrivals.ndim != 1 or np.any(~np.isfinite(arrivals)) or np.any(np.diff(arrivals) < 0)`) and fail fast if violated.
    if arrivals.ndim!=1 or np.any(~np.isfinite(arrivals)) or np.any(np.diff(arrivals)<0):
        raise ValueError('arrivals must be a finite ordered vector')
    # Guard input contract (`not np.isfinite(service_ms) or service_ms <= 0`) and fail fast if violated.
    if not np.isfinite(service_ms) or service_ms<=0:
        raise ValueError('service duration must be positive milliseconds')
    ready=0.
    starts=[]
    finishes=[]
    for arrival in arrivals:
        start=max(float(arrival),ready)
        ready=start+service_ms
        starts.append(start)
        finishes.append(ready)
    return np.array(starts),np.array(finishes),np.array(finishes)-arrivals
hand_start,hand_finish,hand_latency=simulate_queue([0.,1.,4.],2.)
# Execute `np.testing.assert_allclose(hand_start,[0.,2.,4.])`
np.testing.assert_allclose(hand_start,[0.,2.,4.])
# Execute `np.testing.assert_allclose(hand_finish,[2.,4.,6.])`
np.testing.assert_allclose(hand_finish,[2.,4.,6.])
# Execute `np.testing.assert_allclose(hand_latency,[2.,3.,2.])`
np.testing.assert_allclose(hand_latency,[2.,3.,2.])
request_index=np.arange(60)
stable_arrivals=request_index*1.5*service_ms
overloaded_arrivals=request_index*.6*service_ms
_,_,stable_latency=simulate_queue(stable_arrivals,service_ms)
_,_,overloaded_latency=simulate_queue(overloaded_arrivals,service_ms)
assert np.allclose(stable_latency,service_ms)
assert np.isclose(overloaded_latency[-1],24.6*service_ms)
print('measured batch-one service proxy (ms):',service_ms)
print('simulated last stable/overloaded response (ms):',stable_latency[-1],overloaded_latency[-1])

# Step 1 — 1. Define and check the inference workload: The timer includes dispatch and completion for already resident...
import math
import time
import numpy as np
import jax
import jax.numpy as jnp
# Draw pseudorandom samples for `rng` using the explicit RNG state.
rng=np.random.default_rng(82)
# Create device-backed JAX array `W1`.
W1=jnp.asarray(rng.normal(0,.1,(64,32)).astype(np.float32))
# Create device-backed JAX array `W2`.
W2=jnp.asarray(rng.normal(0,.1,(32,8)).astype(np.float32))
# Define and JIT-compile `infer(inputs)` so XLA traces and fuses the operations:
@jax.jit
def infer(inputs):
    return jnp.tanh(inputs@W1)@W2
def measure(batch_size,repeats=40):
    host=rng.normal(size=(batch_size,64)).astype(np.float32)
    # Create device-backed JAX array `device`.
    device=jnp.asarray(host)
    began=time.perf_counter()
    # Synchronize host execution until asynchronous device computation completes.
    first=infer(device).block_until_ready()
    first_ms=(time.perf_counter()-began)*1000
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected=np.tanh(host@np.asarray(W1))@np.asarray(W2)
    np.testing.assert_allclose(first,expected,rtol=2e-5,atol=2e-6)
    samples=[]
    for _ in range(repeats):
        began=time.perf_counter()
        # Synchronize host execution until asynchronous device computation completes.
        infer(device).block_until_ready()
        samples.append((time.perf_counter()-began)*1000)
    p50=float(np.percentile(samples,50))
    return {'batch':batch_size,'first_call_ms':first_ms,'samples_ms':samples,
            'p50_ms':p50,'p95_ms':float(np.percentile(samples,95)),
            'steady_examples_per_second':1000*batch_size/p50}

# Step 2 — 2. Measure each batch shape separately: Do not assume a larger batch is faster per request.
measurements=[measure(batch) for batch in (1,8,32)]
for report in measurements:
    assert len(report['samples_ms'])==40
    # Ensure all array elements remain finite: `all(value>0 and np.isfinite(value) for value in report['samples_m...`
    assert all(value>0 and np.isfinite(value) for value in report['samples_ms'])
    print({key:value for key,value in report.items() if key!='samples_ms'})
service_ms=measurements[0]['p50_ms']

# Step 3 — 3. Model a queue with an independent hand-check: This simulator assumes one serial worker and a constant service...
def simulate_queue(arrival_ms,service_ms):
    # Convert `arrivals` to a host NumPy array for inspection or verification.
    arrivals=np.asarray(arrival_ms,dtype=float)
    # Guard input contract (`arrivals.ndim != 1 or np.any(~np.isfinite(arrivals)) or np.any(np.diff(arrivals) < 0)`) and fail fast if violated.
    if arrivals.ndim!=1 or np.any(~np.isfinite(arrivals)) or np.any(np.diff(arrivals)<0):
        raise ValueError('arrivals must be a finite ordered vector')
    # Guard input contract (`not np.isfinite(service_ms) or service_ms <= 0`) and fail fast if violated.
    if not np.isfinite(service_ms) or service_ms<=0:
        raise ValueError('service duration must be positive milliseconds')
    ready=0.
    starts=[]
    finishes=[]
    for arrival in arrivals:
        start=max(float(arrival),ready)
        ready=start+service_ms
        starts.append(start)
        finishes.append(ready)
    return np.array(starts),np.array(finishes),np.array(finishes)-arrivals
hand_start,hand_finish,hand_latency=simulate_queue([0.,1.,4.],2.)
# Execute `np.testing.assert_allclose(hand_start,[0.,2.,4.])`
np.testing.assert_allclose(hand_start,[0.,2.,4.])
# Execute `np.testing.assert_allclose(hand_finish,[2.,4.,6.])`
np.testing.assert_allclose(hand_finish,[2.,4.,6.])
# Execute `np.testing.assert_allclose(hand_latency,[2.,3.,2.])`
np.testing.assert_allclose(hand_latency,[2.,3.,2.])
request_index=np.arange(60)
stable_arrivals=request_index*1.5*service_ms
overloaded_arrivals=request_index*.6*service_ms
_,_,stable_latency=simulate_queue(stable_arrivals,service_ms)
_,_,overloaded_latency=simulate_queue(overloaded_arrivals,service_ms)
assert np.allclose(stable_latency,service_ms)
assert np.isclose(overloaded_latency[-1],24.6*service_ms)
print('measured batch-one service proxy (ms):',service_ms)
print('simulated last stable/overloaded response (ms):',stable_latency[-1],overloaded_latency[-1])

# Figure data experiment
# Compute figure data for: Measured batch computation and a separately simulated queue
visual_data={"kind":"panels","panels":[{"kind":"bar","x":[0,1,2],"labels":["batch 1","batch 8","batch 32"],"xlabel":"resident-input batch shape","ylabel":"measured CPU milliseconds per batch","series":[{"label":"warm p50","y":[r["p50_ms"] for r in measurements]},{"label":"warm p95","y":[r["p95_ms"] for r in measurements]}]},{"kind":"line","x":request_index.tolist(),"xlabel":"request index (simulation)","ylabel":"simulated response / measured service","series":[{"label":"spaced arrivals","y":(stable_latency/service_ms).tolist()},{"label":"overloaded arrivals","y":(overloaded_latency/service_ms).tolist()}]}]}

# Experiment: Compute a replica estimate with explicit units
# Experiment — Compute a replica estimate with explicit units: This is arithmetic on hypothetical effective capacities, not an...
arrival_per_s=300.
capacity_per_s=200.
target_utilization=.7
replicas=math.ceil(arrival_per_s/(target_utilization*capacity_per_s))
assert replicas==3
print("hypothetical planned replicas:",replicas)

# Experiment: Estimate batch-collection delay
# Experiment — Estimate batch-collection delay: A throughput improvement can be overwhelmed by waiting for a...
batch=8
spacing_ms=2.
collection_wait=(batch-1-np.arange(batch))*spacing_ms
assert collection_wait[0]==14.
assert collection_wait.mean()==7.

# Experiment: Inspect the interpolation behind p95
# Experiment — Inspect the interpolation behind p95: These deliberately constructed values explain the percentile...
illustrative_ms=np.array([1.]*19+[100.])
interpolated_p95=float(np.percentile(illustrative_ms,95,method='linear'))
np.testing.assert_allclose(interpolated_p95,5.95,atol=1e-10)
assert interpolated_p95 not in illustrative_ms and illustrative_ms.max()==100.
print('Analytic sample p95 / max (ms):',interpolated_p95,illustrative_ms.max())

# Reference solution. Try the exercise before reading this.
# Exercise solution: Simulate a burst of four requests arriving at time zero with a...
_,burst_finish,burst_latency=simulate_queue([0.,0.,0.,0.],2.)
# Execute `np.testing.assert_allclose(burst_finish,[2.,4.,6.,8.])`
np.testing.assert_allclose(burst_finish,[2.,4.,6.,8.])
# Execute `np.testing.assert_allclose(burst_latency,[2.,4.,6.,8.])`
np.testing.assert_allclose(burst_latency,[2.,4.,6.,8.])

# Reference practice: Expose burst sensitivity
# Expose burst sensitivity (Transfer / diagnosis): An average arrival rate alone does not describe a burst or...
_,_,smooth=simulate_queue(np.arange(5)*3.,2.)
# Allocate initialized array `(_, _, burst)` with the specified shape and dtype.
_,_,burst=simulate_queue(np.zeros(5),2.)
assert smooth.max()==2. and burst.max()==10.

# Reference practice: Reject a broken timeline
# Reject a broken timeline (Transfer / diagnosis): Validating assumptions prevents a simulation from silently...
for arrivals,duration in [([2.,1.],2.),([0.,1.],0.)]:
    try:
        simulate_queue(arrivals,duration)
    except ValueError:
        pass
    else:
        raise AssertionError("invalid timeline was accepted")

# Reference practice: Measure the margin below a response deadline
# Measure the margin below a response deadline (Transfer / diagnosis): Two responses miss the deadline in this serial simulation.
_,_,deadline_latencies=simulate_queue(np.zeros(5),2.)
# Execute `np.testing.assert_array_equal(deadline_latencies,[2.,4.,6.,8`
np.testing.assert_array_equal(deadline_latencies,[2.,4.,6.,8.,10.])
missed_deadlines=int(np.count_nonzero(deadline_latencies>6.))
assert missed_deadlines==2
print('Requests beyond the declared deadline:',missed_deadlines)
print("PASS: deployment-04")